From 64991456c1f30805ccaffdfeda59e6b3757f5d06 Mon Sep 17 00:00:00 2001 From: Daphne Hansell <128793799+daphnehanse11@users.noreply.github.com> Date: Wed, 18 Mar 2026 15:26:28 -0400 Subject: [PATCH 01/31] Add 2025 post-calibration ACA takeup override --- .../calibration/publish_local_area.py | 47 +++++++++++ .../test_stacked_dataset_builder.py | 17 ++++ .../test_unified_calibration.py | 41 +++++++++- policyengine_us_data/utils/takeup.py | 82 ++++++++++++++++--- 4 files changed, 176 insertions(+), 11 deletions(-) diff --git a/policyengine_us_data/calibration/publish_local_area.py b/policyengine_us_data/calibration/publish_local_area.py index 72594631e..35afbdeb2 100644 --- a/policyengine_us_data/calibration/publish_local_area.py +++ b/policyengine_us_data/calibration/publish_local_area.py @@ -32,8 +32,11 @@ assign_random_geography, ) from policyengine_us_data.utils.takeup import ( + ACA_POST_CALIBRATION_PERSON_TARGETS, SIMPLE_TAKEUP_VARS, apply_block_takeup_to_arrays, + compute_block_takeup_draws_for_entities, + extend_aca_takeup_to_match_target, ) CHECKPOINT_FILE = Path("completed_states.txt") @@ -536,7 +539,51 @@ def build_h5( time_period=time_period, takeup_filter=takeup_filter, ) + + if ( + (takeup_filter is None or "takes_up_aca_if_eligible" in takeup_filter) + and "takes_up_aca_if_eligible" in takeup_results + and 2025 in ACA_POST_CALIBRATION_PERSON_TARGETS + ): + print("Applying 2025 ACA post-calibration enrollment override...") + sim.set_input( + "takes_up_aca_if_eligible", + 2025, + np.ones(len(entity_id_arrays["tax_unit"]), dtype=bool), + ) + sim.delete_arrays("aca_ptc") + enrolled_if_takeup = ( + sim.calculate("aca_ptc", map_to="person", period=2025).values > 0 + )[person_clone_idx] + aca_person_weights_if_takeup = np.zeros( + len(entity_clone_idx["tax_unit"]), + dtype=np.float64, + ) + np.add.at( + aca_person_weights_if_takeup, + new_person_entity_ids["tax_unit"], + enrolled_if_takeup.astype(np.float64) * person_weights, + ) + tax_unit_hh_idx = entity_hh_indices["tax_unit"] + aca_draws = compute_block_takeup_draws_for_entities( + "takes_up_aca_if_eligible", + active_blocks[tax_unit_hh_idx].astype(str), + original_hh_ids[tax_unit_hh_idx], + entity_clone_ids=tax_unit_hh_idx, + ) + data["takes_up_aca_if_eligible"] = { + time_period: takeup_results["takes_up_aca_if_eligible"], + 2025: extend_aca_takeup_to_match_target( + takeup_results["takes_up_aca_if_eligible"], + aca_draws, + aca_person_weights_if_takeup, + ACA_POST_CALIBRATION_PERSON_TARGETS[2025], + ), + } + for var_name, bools in takeup_results.items(): + if var_name == "takes_up_aca_if_eligible" and var_name in data: + continue data[var_name] = {time_period: bools} # === Write H5 === diff --git a/policyengine_us_data/tests/test_calibration/test_stacked_dataset_builder.py b/policyengine_us_data/tests/test_calibration/test_stacked_dataset_builder.py index 339dec4e6..df88d70ce 100644 --- a/policyengine_us_data/tests/test_calibration/test_stacked_dataset_builder.py +++ b/policyengine_us_data/tests/test_calibration/test_stacked_dataset_builder.py @@ -167,6 +167,23 @@ def test_household_count_matches_weights(self, stacked_result, test_weights): assert len(hh_df) == expected_households +class TestAcaTakeupOverride: + """Verify stacked datasets carry a 2025 ACA takeup override.""" + + def test_aca_takeup_includes_2025_period(self, stacked_sim): + data = stacked_sim.dataset.load_dataset()["takes_up_aca_if_eligible"] + periods = {int(period) for period in data} + assert periods == {int(stacked_sim.default_calculation_period), 2025} + + def test_aca_takeup_2025_only_adds_true_values(self, stacked_sim): + data = stacked_sim.dataset.load_dataset()["takes_up_aca_if_eligible"] + base_period = int(stacked_sim.default_calculation_period) + takeup_2024 = np.asarray(data[next(p for p in data if int(p) == base_period)]) + takeup_2025 = np.asarray(data[next(p for p in data if int(p) == 2025)]) + assert takeup_2025.mean() > takeup_2024.mean() + assert np.all(takeup_2024 <= takeup_2025) + + @pytest.fixture(scope="module") def stacked_sim(test_weights, n_households): """Run stacked dataset builder and return the simulation.""" diff --git a/policyengine_us_data/tests/test_calibration/test_unified_calibration.py b/policyengine_us_data/tests/test_calibration/test_unified_calibration.py index 28a3c906f..ca30d2980 100644 --- a/policyengine_us_data/tests/test_calibration/test_unified_calibration.py +++ b/policyengine_us_data/tests/test_calibration/test_unified_calibration.py @@ -12,8 +12,10 @@ from policyengine_us_data.utils.takeup import ( SIMPLE_TAKEUP_VARS, TAKEUP_AFFECTED_TARGETS, - compute_block_takeup_for_entities, apply_block_takeup_to_arrays, + compute_block_takeup_draws_for_entities, + compute_block_takeup_for_entities, + extend_aca_takeup_to_match_target, _resolve_rate, ) from policyengine_us_data.calibration.clone_and_assign import ( @@ -178,6 +180,43 @@ def test_different_blocks_different_result(self): assert differs +class TestAcaTakeupTargeting: + """Verify ACA post-calibration targeting helpers.""" + + def test_draw_helper_matches_boolean_helper(self): + blocks = np.array(["370010001001001"] * 25) + hh_ids = np.arange(25, dtype=np.int64) + draws = compute_block_takeup_draws_for_entities( + "takes_up_aca_if_eligible", + blocks, + hh_ids, + ) + result = compute_block_takeup_for_entities( + "takes_up_aca_if_eligible", + 0.7, + blocks, + hh_ids, + ) + np.testing.assert_array_equal(result, draws < 0.7) + + def test_extend_only_adds_true_values_until_target(self): + base_takeup = np.array([True, False, False, False], dtype=bool) + entity_draws = np.array([0.10, 0.40, 0.20, 0.30], dtype=np.float64) + enrolled_person_weights = np.array([2.0, 1.0, 3.0, 4.0], dtype=np.float64) + + result = extend_aca_takeup_to_match_target( + base_takeup, + entity_draws, + enrolled_person_weights, + target_people=6.0, + ) + + np.testing.assert_array_equal( + result, + np.array([True, False, True, True], dtype=bool), + ) + + class TestResolveRate: """Verify _resolve_rate handles scalar and dict rates.""" diff --git a/policyengine_us_data/utils/takeup.py b/policyengine_us_data/utils/takeup.py index 5e49b20ac..dc35376ed 100644 --- a/policyengine_us_data/utils/takeup.py +++ b/policyengine_us_data/utils/takeup.py @@ -108,6 +108,12 @@ }, } +# CMS 2025 Marketplace OEP State-Level Public Use File, Total / All row. +# This is the number of consumers receiving APTC in plan year 2025. +ACA_POST_CALIBRATION_PERSON_TARGETS = { + 2025: 22_380_137, +} + # FIPS -> 2-letter state code for Medicaid rate lookup _FIPS_TO_STATE_CODE = { 1: "AL", @@ -178,14 +184,13 @@ def _resolve_rate( return float(rate_or_dict) -def compute_block_takeup_for_entities( +def compute_block_takeup_draws_for_entities( var_name: str, - rate_or_dict, entity_blocks: np.ndarray, entity_hh_ids: np.ndarray = None, entity_clone_ids: np.ndarray = None, ) -> np.ndarray: - """Compute boolean takeup via block-level seeded draws. + """Compute deterministic uniform draws for block-level takeup. Each unique (block, household) pair gets its own seeded RNG, producing reproducible draws regardless of how many households @@ -197,11 +202,10 @@ def compute_block_takeup_for_entities( builder, which processes each clone independently. State FIPS for rate resolution is derived from the first two - characters of each block GEOID. + characters of each block GEOID in the downstream boolean step. Args: var_name: Takeup variable name. - rate_or_dict: Scalar rate or {state_code: rate} dict. entity_blocks: Block GEOID per entity (str array). entity_hh_ids: Household ID per entity (int array). When provided, seeds per (block, household) for @@ -212,19 +216,15 @@ def compute_block_takeup_for_entities( takeup decisions. Returns: - Boolean array of shape (n_entities,). + Float array of shape (n_entities,) in [0, 1). """ n = len(entity_blocks) draws = np.zeros(n, dtype=np.float64) - rates = np.ones(n, dtype=np.float64) for block in np.unique(entity_blocks): if block == "": continue blk_mask = entity_blocks == block - sf = int(str(block)[:2]) - rate = _resolve_rate(rate_or_dict, sf) - rates[blk_mask] = rate if entity_hh_ids is not None: for hh_id in np.unique(entity_hh_ids[blk_mask]): @@ -248,9 +248,71 @@ def compute_block_takeup_for_entities( rng = seeded_rng(var_name, salt=str(block)) draws[blk_mask] = rng.random(int(blk_mask.sum())) + return draws + + +def compute_block_takeup_for_entities( + var_name: str, + rate_or_dict, + entity_blocks: np.ndarray, + entity_hh_ids: np.ndarray = None, + entity_clone_ids: np.ndarray = None, +) -> np.ndarray: + """Compute boolean takeup via block-level seeded draws.""" + draws = compute_block_takeup_draws_for_entities( + var_name, + entity_blocks, + entity_hh_ids, + entity_clone_ids, + ) + rates = np.ones(len(entity_blocks), dtype=np.float64) + + for block in np.unique(entity_blocks): + if block == "": + continue + blk_mask = entity_blocks == block + rates[blk_mask] = _resolve_rate(rate_or_dict, int(str(block)[:2])) + return draws < rates +def extend_aca_takeup_to_match_target( + base_takeup: np.ndarray, + entity_draws: np.ndarray, + enrolled_person_weights: np.ndarray, + target_people: float, +) -> np.ndarray: + """Turn on extra tax units until weighted ACA enrollment hits target. + + ``enrolled_person_weights`` should be the weighted number of + people in each tax unit who would receive ACA PTC if that tax + unit takes up coverage in the target year. + """ + result = base_takeup.copy() + current_people = enrolled_person_weights[result].sum() + if current_people >= target_people: + return result + + available_mask = (~result) & (enrolled_person_weights > 0) + if not available_mask.any(): + return result + + available_idx = np.flatnonzero(available_mask) + ordered_idx = available_idx[ + np.argsort(entity_draws[available_idx], kind="stable") + ] + cumulative_people = current_people + np.cumsum( + enrolled_person_weights[ordered_idx] + ) + n_to_add = np.searchsorted( + cumulative_people, + target_people, + side="left", + ) + 1 + result[ordered_idx[:n_to_add]] = True + return result + + def apply_block_takeup_to_arrays( hh_blocks: np.ndarray, hh_state_fips: np.ndarray, From 04acc8be3d8499da2fc4111c865489c1610b8da8 Mon Sep 17 00:00:00 2001 From: Daphne Hansell <128793799+daphnehanse11@users.noreply.github.com> Date: Wed, 18 Mar 2026 15:34:33 -0400 Subject: [PATCH 02/31] Fix lint in ACA takeup tests --- .../tests/test_calibration/test_unified_calibration.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/policyengine_us_data/tests/test_calibration/test_unified_calibration.py b/policyengine_us_data/tests/test_calibration/test_unified_calibration.py index ca30d2980..382354cc1 100644 --- a/policyengine_us_data/tests/test_calibration/test_unified_calibration.py +++ b/policyengine_us_data/tests/test_calibration/test_unified_calibration.py @@ -6,8 +6,6 @@ """ import numpy as np -import pytest - from policyengine_us_data.utils.randomness import seeded_rng from policyengine_us_data.utils.takeup import ( SIMPLE_TAKEUP_VARS, @@ -556,7 +554,6 @@ def test_matrix_and_stacked_identical_draws(self): def test_aggregation_entity_to_household(self): """np.add.at aggregation matches manual per-HH sum.""" n_hh = 3 - n_ent = 6 ent_hh = np.array([0, 0, 1, 1, 1, 2]) eligible = np.array( [100.0, 200.0, 50.0, 150.0, 100.0, 300.0], From e994f017ad2cd734708be8827bf64a66f08ffd60 Mon Sep 17 00:00:00 2001 From: Daphne Hansell <128793799+daphnehanse11@users.noreply.github.com> Date: Wed, 18 Mar 2026 15:41:54 -0400 Subject: [PATCH 03/31] Format ACA takeup helper --- policyengine_us_data/utils/takeup.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/policyengine_us_data/utils/takeup.py b/policyengine_us_data/utils/takeup.py index dc35376ed..34790ed06 100644 --- a/policyengine_us_data/utils/takeup.py +++ b/policyengine_us_data/utils/takeup.py @@ -298,17 +298,16 @@ def extend_aca_takeup_to_match_target( return result available_idx = np.flatnonzero(available_mask) - ordered_idx = available_idx[ - np.argsort(entity_draws[available_idx], kind="stable") - ] - cumulative_people = current_people + np.cumsum( - enrolled_person_weights[ordered_idx] + ordered_idx = available_idx[np.argsort(entity_draws[available_idx], kind="stable")] + cumulative_people = current_people + np.cumsum(enrolled_person_weights[ordered_idx]) + n_to_add = ( + np.searchsorted( + cumulative_people, + target_people, + side="left", + ) + + 1 ) - n_to_add = np.searchsorted( - cumulative_people, - target_people, - side="left", - ) + 1 result[ordered_idx[:n_to_add]] = True return result From 1b0bd68c6ba0c30ecc4f5ae7e9b4bab8cbc0f445 Mon Sep 17 00:00:00 2001 From: Daphne Hansell <128793799+daphnehanse11@users.noreply.github.com> Date: Wed, 18 Mar 2026 15:55:25 -0400 Subject: [PATCH 04/31] Move ACA override to Enhanced CPS path --- .../calibration/publish_local_area.py | 47 -------- .../datasets/cps/enhanced_cps.py | 110 ++++++++++++++++-- .../test_stacked_dataset_builder.py | 17 --- .../tests/test_datasets/test_enhanced_cps.py | 22 +++- 4 files changed, 121 insertions(+), 75 deletions(-) diff --git a/policyengine_us_data/calibration/publish_local_area.py b/policyengine_us_data/calibration/publish_local_area.py index 35afbdeb2..72594631e 100644 --- a/policyengine_us_data/calibration/publish_local_area.py +++ b/policyengine_us_data/calibration/publish_local_area.py @@ -32,11 +32,8 @@ assign_random_geography, ) from policyengine_us_data.utils.takeup import ( - ACA_POST_CALIBRATION_PERSON_TARGETS, SIMPLE_TAKEUP_VARS, apply_block_takeup_to_arrays, - compute_block_takeup_draws_for_entities, - extend_aca_takeup_to_match_target, ) CHECKPOINT_FILE = Path("completed_states.txt") @@ -539,51 +536,7 @@ def build_h5( time_period=time_period, takeup_filter=takeup_filter, ) - - if ( - (takeup_filter is None or "takes_up_aca_if_eligible" in takeup_filter) - and "takes_up_aca_if_eligible" in takeup_results - and 2025 in ACA_POST_CALIBRATION_PERSON_TARGETS - ): - print("Applying 2025 ACA post-calibration enrollment override...") - sim.set_input( - "takes_up_aca_if_eligible", - 2025, - np.ones(len(entity_id_arrays["tax_unit"]), dtype=bool), - ) - sim.delete_arrays("aca_ptc") - enrolled_if_takeup = ( - sim.calculate("aca_ptc", map_to="person", period=2025).values > 0 - )[person_clone_idx] - aca_person_weights_if_takeup = np.zeros( - len(entity_clone_idx["tax_unit"]), - dtype=np.float64, - ) - np.add.at( - aca_person_weights_if_takeup, - new_person_entity_ids["tax_unit"], - enrolled_if_takeup.astype(np.float64) * person_weights, - ) - tax_unit_hh_idx = entity_hh_indices["tax_unit"] - aca_draws = compute_block_takeup_draws_for_entities( - "takes_up_aca_if_eligible", - active_blocks[tax_unit_hh_idx].astype(str), - original_hh_ids[tax_unit_hh_idx], - entity_clone_ids=tax_unit_hh_idx, - ) - data["takes_up_aca_if_eligible"] = { - time_period: takeup_results["takes_up_aca_if_eligible"], - 2025: extend_aca_takeup_to_match_target( - takeup_results["takes_up_aca_if_eligible"], - aca_draws, - aca_person_weights_if_takeup, - ACA_POST_CALIBRATION_PERSON_TARGETS[2025], - ), - } - for var_name, bools in takeup_results.items(): - if var_name == "takes_up_aca_if_eligible" and var_name in data: - continue data[var_name] = {time_period: bools} # === Write H5 === diff --git a/policyengine_us_data/datasets/cps/enhanced_cps.py b/policyengine_us_data/datasets/cps/enhanced_cps.py index 8755c73e4..0c3d391f0 100644 --- a/policyengine_us_data/datasets/cps/enhanced_cps.py +++ b/policyengine_us_data/datasets/cps/enhanced_cps.py @@ -1,10 +1,7 @@ from policyengine_core.data import Dataset import pandas as pd from policyengine_us_data.utils import ( - pe_to_soi, - get_soi, build_loss_matrix, - fmt, HardConcrete, print_reweighting_diagnostics, set_seeds, @@ -15,10 +12,14 @@ from typing import Type from policyengine_us_data.storage import STORAGE_FOLDER from policyengine_us_data.datasets.cps.extended_cps import ( - ExtendedCPS_2024, ExtendedCPS_2024_Half, CPS_2024, ) +from policyengine_us_data.utils.randomness import seeded_rng +from policyengine_us_data.utils.takeup import ( + ACA_POST_CALIBRATION_PERSON_TARGETS, + extend_aca_takeup_to_match_target, +) import logging try: @@ -27,6 +28,48 @@ torch = None +def _get_period_array(period_values: dict, period: int) -> np.ndarray: + """Get a period array from a TIME_PERIOD_ARRAYS variable dict.""" + value = period_values.get(period) + if value is None: + value = period_values.get(str(period)) + if value is None: + raise KeyError(f"Missing period {period}") + return np.asarray(value) + + +def create_aca_2025_takeup_override( + base_takeup: np.ndarray, + person_enrolled_if_takeup: np.ndarray, + person_weights: np.ndarray, + person_tax_unit_ids: np.ndarray, + tax_unit_ids: np.ndarray, + target_people: float = ACA_POST_CALIBRATION_PERSON_TARGETS[2025], +) -> np.ndarray: + """Add 2025 ACA takers until weighted APTC enrollment hits target.""" + tax_unit_id_to_idx = { + int(tax_unit_id): idx for idx, tax_unit_id in enumerate(tax_unit_ids) + } + person_tax_unit_idx = np.array( + [tax_unit_id_to_idx[int(tax_unit_id)] for tax_unit_id in person_tax_unit_ids], + dtype=np.int64, + ) + enrolled_person_weights = np.zeros(len(tax_unit_ids), dtype=np.float64) + np.add.at( + enrolled_person_weights, + person_tax_unit_idx, + person_enrolled_if_takeup.astype(np.float64) * person_weights, + ) + draws = seeded_rng("takes_up_aca_if_eligible").random(len(tax_unit_ids)) + + return extend_aca_takeup_to_match_target( + base_takeup=np.asarray(base_takeup, dtype=bool), + entity_draws=draws, + enrolled_person_weights=enrolled_person_weights, + target_people=target_people, + ) + + def reweight( original_weights, loss_matrix, @@ -88,7 +131,7 @@ def loss(weights): optimizer.zero_grad() masked = torch.exp(weights) * gates() l_main = loss(masked) - l = l_main + l0_lambda * gates.get_penalty() + loss_value = l_main + l0_lambda * gates.get_penalty() if (log_path is not None) and (i % 10 == 0): gates.eval() estimates = (torch.exp(weights) * gates()) @ loss_matrix @@ -112,10 +155,12 @@ def loss(weights): if (log_path is not None) and (i % 1000 == 0): performance.to_csv(log_path, index=False) if start_loss is None: - start_loss = l.item() - loss_rel_change = (l.item() - start_loss) / start_loss - l.backward() - iterator.set_postfix({"loss": l.item(), "loss_rel_change": loss_rel_change}) + start_loss = loss_value.item() + loss_rel_change = (loss_value.item() - start_loss) / start_loss + loss_value.backward() + iterator.set_postfix( + {"loss": loss_value.item(), "loss_rel_change": loss_rel_change} + ) optimizer.step() if log_path is not None: performance.to_csv(log_path, index=False) @@ -144,6 +189,7 @@ def generate(self): sim = Microsimulation(dataset=self.input_dataset) data = sim.dataset.load_dataset() + base_year = int(sim.default_calculation_period) data["household_weight"] = {} original_weights = sim.calculate("household_weight") original_weights = original_weights.values + np.random.normal( @@ -218,6 +264,52 @@ def generate(self): f"{int(np.sum(w > 0))} non-zero" ) + if 2025 in ACA_POST_CALIBRATION_PERSON_TARGETS: + sim.set_input( + "household_weight", + base_year, + _get_period_array(data["household_weight"], base_year).astype( + np.float32 + ), + ) + sim.set_input( + "takes_up_aca_if_eligible", + 2025, + np.ones( + len(_get_period_array(data["tax_unit_id"], base_year)), + dtype=bool, + ), + ) + sim.delete_arrays("aca_ptc") + + data["takes_up_aca_if_eligible"][2025] = create_aca_2025_takeup_override( + base_takeup=_get_period_array( + data["takes_up_aca_if_eligible"], + base_year, + ), + person_enrolled_if_takeup=np.asarray( + sim.calculate( + "aca_ptc", + map_to="person", + period=2025, + use_weights=False, + ) + ) + > 0, + person_weights=np.asarray( + sim.calculate( + "person_weight", + period=2025, + use_weights=False, + ) + ), + person_tax_unit_ids=_get_period_array( + data["person_tax_unit_id"], + base_year, + ), + tax_unit_ids=_get_period_array(data["tax_unit_id"], base_year), + ) + logging.info("Post-generation weight validation passed") self.save_dataset(data) diff --git a/policyengine_us_data/tests/test_calibration/test_stacked_dataset_builder.py b/policyengine_us_data/tests/test_calibration/test_stacked_dataset_builder.py index df88d70ce..339dec4e6 100644 --- a/policyengine_us_data/tests/test_calibration/test_stacked_dataset_builder.py +++ b/policyengine_us_data/tests/test_calibration/test_stacked_dataset_builder.py @@ -167,23 +167,6 @@ def test_household_count_matches_weights(self, stacked_result, test_weights): assert len(hh_df) == expected_households -class TestAcaTakeupOverride: - """Verify stacked datasets carry a 2025 ACA takeup override.""" - - def test_aca_takeup_includes_2025_period(self, stacked_sim): - data = stacked_sim.dataset.load_dataset()["takes_up_aca_if_eligible"] - periods = {int(period) for period in data} - assert periods == {int(stacked_sim.default_calculation_period), 2025} - - def test_aca_takeup_2025_only_adds_true_values(self, stacked_sim): - data = stacked_sim.dataset.load_dataset()["takes_up_aca_if_eligible"] - base_period = int(stacked_sim.default_calculation_period) - takeup_2024 = np.asarray(data[next(p for p in data if int(p) == base_period)]) - takeup_2025 = np.asarray(data[next(p for p in data if int(p) == 2025)]) - assert takeup_2025.mean() > takeup_2024.mean() - assert np.all(takeup_2024 <= takeup_2025) - - @pytest.fixture(scope="module") def stacked_sim(test_weights, n_households): """Run stacked dataset builder and return the simulation.""" diff --git a/policyengine_us_data/tests/test_datasets/test_enhanced_cps.py b/policyengine_us_data/tests/test_datasets/test_enhanced_cps.py index 298de5a4a..4df2ae410 100644 --- a/policyengine_us_data/tests/test_datasets/test_enhanced_cps.py +++ b/policyengine_us_data/tests/test_datasets/test_enhanced_cps.py @@ -1,4 +1,4 @@ -import pytest +import numpy as np def test_ecps_employment_income_direct(): @@ -97,7 +97,6 @@ def apply(self): def test_ssn_card_type_none_target(): from policyengine_us_data.datasets.cps import EnhancedCPS_2024 from policyengine_us import Microsimulation - import numpy as np TARGET_COUNT = 13e6 TOLERANCE = 0.2 # Allow ±20% error @@ -184,6 +183,25 @@ def test_aca_calibration(): assert not failed, f"One or more states exceeded tolerance of {TOLERANCE:.0%}." +def test_aca_2025_takeup_override_helper(): + from policyengine_us_data.datasets.cps.enhanced_cps import ( + create_aca_2025_takeup_override, + ) + + result = create_aca_2025_takeup_override( + base_takeup=np.array([True, False, False], dtype=bool), + person_enrolled_if_takeup=np.array([True, True, True, True], dtype=bool), + person_weights=np.array([2.0, 1.0, 3.0, 4.0], dtype=np.float64), + person_tax_unit_ids=np.array([10, 10, 11, 12], dtype=np.int64), + tax_unit_ids=np.array([10, 11, 12], dtype=np.int64), + target_people=6.0, + ) + + assert np.all(np.array([True, False, False]) <= result) + assert result.dtype == bool + assert result.sum() == 2 + + def test_immigration_status_diversity(): """Test that immigration statuses show appropriate diversity (not all citizens).""" from policyengine_us_data.datasets.cps import EnhancedCPS_2024 From 24efc95c5be2054772ce9013e1ebf0a618d8a952 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 30 Mar 2026 22:23:26 -0400 Subject: [PATCH 05/31] Add long-run calibration contracts --- .../datasets/cps/long_term/README.md | 51 +- .../long_term/assess_calibration_frontier.py | 247 +++++++++ .../datasets/cps/long_term/calibration.py | 496 +++++++++++++++++- .../cps/long_term/calibration_artifacts.py | 202 +++++++ .../cps/long_term/calibration_profiles.py | 379 +++++++++++++ .../long_term/rebuild_calibration_manifest.py | 40 ++ .../cps/long_term/run_household_projection.py | 196 +++++-- .../datasets/cps/long_term/ssa_data.py | 114 ++-- policyengine_us_data/storage/README.md | 7 + .../long_term_target_sources/sources.json | 18 + .../trustees_2025_current_law.csv | 77 +++ .../test_long_term_calibration_contract.py | 390 ++++++++++++++ 12 files changed, 2137 insertions(+), 80 deletions(-) create mode 100644 policyengine_us_data/datasets/cps/long_term/assess_calibration_frontier.py create mode 100644 policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py create mode 100644 policyengine_us_data/datasets/cps/long_term/calibration_profiles.py create mode 100644 policyengine_us_data/datasets/cps/long_term/rebuild_calibration_manifest.py create mode 100644 policyengine_us_data/storage/long_term_target_sources/sources.json create mode 100644 policyengine_us_data/storage/long_term_target_sources/trustees_2025_current_law.csv create mode 100644 policyengine_us_data/tests/test_long_term_calibration_contract.py diff --git a/policyengine_us_data/datasets/cps/long_term/README.md b/policyengine_us_data/datasets/cps/long_term/README.md index c20216dc8..e8de8dfe0 100644 --- a/policyengine_us_data/datasets/cps/long_term/README.md +++ b/policyengine_us_data/datasets/cps/long_term/README.md @@ -6,24 +6,43 @@ Run projections using `run_household_projection.py`: ```bash -# Recommended: GREG with all constraint types -python run_household_projection.py 2100 --greg --use-ss --use-payroll --use-tob --save-h5 +# Recommended: named profile with all constraint types +python run_household_projection.py 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --save-h5 # IPF with only age distribution constraints (faster, less accurate) -python run_household_projection.py 2050 +python run_household_projection.py 2050 --profile age-only # GREG with age + Social Security only -python run_household_projection.py 2100 --greg --use-ss +python run_household_projection.py 2100 --profile ss ``` **Arguments:** - `END_YEAR`: Target year for projection (default: 2035) +- `--profile`: Named calibration contract. Recommended over legacy flags. +- `--target-source`: Named long-term target source package. +- `--output-dir`: Output directory for generated H5 files and metadata sidecars. - `--greg`: Use GREG calibration instead of IPF - `--use-ss`: Include Social Security benefit totals as calibration target (requires `--greg`) - `--use-payroll`: Include taxable payroll totals as calibration target (requires `--greg`) - `--use-tob`: Include TOB (Taxation of Benefits) revenue as calibration target (requires `--greg`) - `--save-h5`: Save year-specific .h5 files to `./projected_datasets/` directory +**Named profiles:** +- `age-only`: IPF age-only calibration +- `ss`: positive entropy calibration with age + Social Security +- `ss-payroll`: positive entropy calibration with age + Social Security + taxable payroll +- `ss-payroll-tob`: positive entropy calibration with age + Social Security + taxable payroll + TOB +- `ss-payroll-tob-h6`: positive entropy calibration with age + Social Security + taxable payroll + TOB + H6 + +**Validation contract:** +- Economic-targeted profiles no longer silently pretend an IPF fallback is equivalent to GREG. +- Named economic profiles must produce non-negative weights. +- Each generated H5 now gets a `YYYY.h5.metadata.json` sidecar with profile and calibration audit details. +- Each generated H5 sidecar now records the named long-term target source used for the build. +- Each output directory now gets a `calibration_manifest.json` file describing the + profile/base dataset contract for the full artifact set. +- Profiles validate achieved constraint errors before writing output. + **Estimated runtime:** ~2 minutes/year without `--save-h5`, ~3 minutes/year with `--save-h5` --- @@ -35,11 +54,16 @@ python run_household_projection.py 2100 --greg --use-ss - Fast and simple, but cannot enforce Social Security or payroll totals - Converges iteratively (typically 20-40 iterations) -**GREG (Generalized Regression Estimator)** -- Solves for weights matching multiple constraints simultaneously +**Positive Entropy Calibration** +- Solves for strictly positive weights matching multiple constraints simultaneously - Can enforce age distribution + Social Security benefits + taxable payroll + TOB revenue -- One-shot solution using `samplics` package -- **Recommended** for accurate long-term projections +- Uses dual optimization to minimize divergence from baseline weights +- **Recommended** for publishable long-term projections + +**GREG (Generalized Regression Estimator)** +- Legacy linear calibration path retained for explicit flag-based runs +- Can hit constraints exactly, but may produce negative weights in far-horizon years +- No longer the default for named economic calibration profiles --- @@ -51,18 +75,18 @@ python run_household_projection.py 2100 --greg --use-ss 2. **Social Security Benefits** (`--use-ss`, GREG only) - Total OASDI benefit payments (nominal dollars) - - Source: SSA Trustee Report 2024 (`social_security_aux.csv`) + - Source: selected long-term target source package 3. **Taxable Payroll** (`--use-payroll`, GREG only) - W-2 wages capped at wage base + SE income within remaining cap room - Calculated as: `taxable_earnings_for_social_security` + `social_security_taxable_self_employment_income` - - Source: SSA Trustee Report 2024 (`social_security_aux.csv`) + - Source: selected long-term target source package 4. **TOB Revenue** (`--use-tob`, GREG only) - Taxation of Benefits revenue for OASDI and Medicare HI trust funds - OASDI: `tob_revenue_oasdi` (tier 1 taxation, 0-50% of benefits) - HI: `tob_revenue_medicare_hi` (tier 2 taxation, 50-85% of benefits) - - Source: SSA Trustee Report 2024 (`social_security_aux.csv`) + - Source: selected long-term target source package --- @@ -80,7 +104,8 @@ python run_household_projection.py 2100 --greg --use-ss **Local files** (in `policyengine_us_data/storage/`): - `SSPopJul_TR2024.csv` - Population projections 2025-2100 by single year of age -- `social_security_aux.csv` - OASDI costs, taxable payroll, and TOB revenue projections 2025-2100 +- `long_term_target_sources/trustees_2025_current_law.csv` - explicit frozen Trustees/current-law package +- `long_term_target_sources/sources.json` - provenance metadata for named source packages --- @@ -102,7 +127,7 @@ python run_household_projection.py 2100 --greg --use-ss - **`run_household_projection.py`** - Main projection script (see Quick Start) - **`calibration.py`** - IPF and GREG weight calibration implementations -- **`ssa_data.py`** - Load SSA population, benefit, and payroll projections +- **`ssa_data.py`** - Load SSA population and named long-term target source projections - **`projection_utils.py`** - Utility functions (age matrix builder, H5 file creator) - **`extract_ssa_costs.py`** - One-time script to extract SSA data from Excel (already run) diff --git a/policyengine_us_data/datasets/cps/long_term/assess_calibration_frontier.py b/policyengine_us_data/datasets/cps/long_term/assess_calibration_frontier.py new file mode 100644 index 000000000..63e6ac627 --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/assess_calibration_frontier.py @@ -0,0 +1,247 @@ +from __future__ import annotations + +import argparse +import csv +import gc +from pathlib import Path + +import numpy as np +from policyengine_us import Microsimulation + +from calibration import ( + _build_constraint_dataframe_and_controls, + assess_nonnegative_feasibility, +) +from calibration_profiles import get_profile +from projection_utils import build_household_age_matrix +from ssa_data import ( + load_hi_tob_projections, + load_oasdi_tob_projections, + load_ssa_age_projections, + load_ssa_benefit_projections, + load_taxable_payroll_projections, +) + + +DATASET_OPTIONS = { + "enhanced_cps_2024": { + "path": "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5", + "base_year": 2024, + }, +} +SELECTED_DATASET = "enhanced_cps_2024" +BASE_DATASET_PATH = DATASET_OPTIONS[SELECTED_DATASET]["path"] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Assess the nonnegative calibration frontier for a named long-term " + "profile by solving the minimax relative-error LP." + ) + ) + parser.add_argument( + "--profile", + default="ss-payroll-tob", + help="Named calibration profile to assess.", + ) + parser.add_argument( + "--start-year", + type=int, + default=2035, + help="First year to assess when --years is not provided.", + ) + parser.add_argument( + "--end-year", + type=int, + default=2100, + help="Last year to assess when --years is not provided.", + ) + parser.add_argument( + "--step", + type=int, + default=5, + help="Year increment when --years is not provided.", + ) + parser.add_argument( + "--years", + help="Optional comma-separated list of explicit years to assess.", + ) + parser.add_argument( + "--output", + help="Optional CSV path for the frontier table.", + ) + return parser.parse_args() + + +def parse_years(args: argparse.Namespace) -> list[int]: + if args.years: + return [int(value.strip()) for value in args.years.split(",") if value.strip()] + return list(range(args.start_year, args.end_year + 1, args.step)) + + +def reorder_to_households(values, order, n_households: int) -> np.ndarray: + ordered = np.zeros(n_households, dtype=float) + ordered[order] = np.asarray(values, dtype=float) + return ordered + + +def build_constraint_inputs(year: int, hh_id_to_idx: dict, n_households: int, profile) -> dict: + sim = Microsimulation(dataset=BASE_DATASET_PATH) + if profile.use_h6_reform: + raise NotImplementedError( + "Frontier assessment for H6-enabled profiles is not yet implemented." + ) + household_ids = sim.calculate("household_id", period=year, map_to="household").values + if len(household_ids) != n_households: + raise ValueError( + f"Household count mismatch for {year}: {len(household_ids)} vs {n_households}" + ) + order = np.fromiter( + (hh_id_to_idx[hh_id] for hh_id in household_ids), + dtype=int, + count=len(household_ids), + ) + inputs: dict[str, np.ndarray | float | None] = { + "ss_values": None, + "ss_target": None, + "payroll_values": None, + "payroll_target": None, + "h6_income_values": None, + "h6_revenue_target": None, + "oasdi_tob_values": None, + "oasdi_tob_target": None, + "hi_tob_values": None, + "hi_tob_target": None, + } + + if profile.use_ss: + inputs["ss_values"] = reorder_to_households( + sim.calculate("social_security", period=year, map_to="household").values, + order, + n_households, + ) + inputs["ss_target"] = load_ssa_benefit_projections(year) + + if profile.use_payroll: + inputs["payroll_values"] = reorder_to_households( + sim.calculate( + "taxable_earnings_for_social_security", + period=year, + map_to="household", + ).values + + sim.calculate( + "social_security_taxable_self_employment_income", + period=year, + map_to="household", + ).values, + order, + n_households, + ) + inputs["payroll_target"] = load_taxable_payroll_projections(year) + + if profile.use_tob: + inputs["oasdi_tob_values"] = reorder_to_households( + sim.calculate( + "tob_revenue_oasdi", + period=year, + map_to="household", + ).values, + order, + n_households, + ) + inputs["hi_tob_values"] = reorder_to_households( + sim.calculate( + "tob_revenue_medicare_hi", + period=year, + map_to="household", + ).values, + order, + n_households, + ) + inputs["oasdi_tob_target"] = load_oasdi_tob_projections(year) + inputs["hi_tob_target"] = load_hi_tob_projections(year) + + del sim + gc.collect() + return inputs + + +def main() -> int: + args = parse_args() + years = parse_years(args) + if not years: + raise ValueError("No years requested.") + + profile = get_profile(args.profile) + start_year = min(years) + end_year = max(years) + target_matrix = load_ssa_age_projections(start_year=start_year, end_year=end_year) + + base_sim = Microsimulation(dataset=BASE_DATASET_PATH) + X, household_ids_unique, hh_id_to_idx = build_household_age_matrix( + base_sim, + n_ages=target_matrix.shape[0], + ) + del base_sim + gc.collect() + + rows: list[dict[str, object]] = [] + print( + f"Assessing profile {profile.name!r} for {len(years)} years " + f"using {len(household_ids_unique):,} fixed households." + ) + for year in years: + year_idx = year - start_year + y_target = target_matrix[:, year_idx] + inputs = build_constraint_inputs( + year, + hh_id_to_idx, + len(household_ids_unique), + profile, + ) + aux_df, controls = _build_constraint_dataframe_and_controls( + X, + y_target, + n_ages=target_matrix.shape[0], + **inputs, + ) + targets = np.array(list(controls.values()), dtype=float) + feasibility = assess_nonnegative_feasibility( + aux_df.to_numpy(dtype=float), + targets, + ) + best_case = feasibility["best_case_max_pct_error"] + within_tolerance = ( + best_case is not None and best_case <= profile.max_constraint_error_pct + ) + row = { + "year": year, + "profile": profile.name, + "best_case_max_pct_error": best_case, + "within_profile_tolerance": within_tolerance, + "quality": "exact" if within_tolerance else "approximate", + "status": feasibility["status"], + "message": feasibility["message"], + } + rows.append(row) + best_case_display = "n/a" if best_case is None else f"{best_case:.3f}%" + print( + f"{year}: best-case max error {best_case_display} -> " + f"{row['quality']}" + ) + + if args.output: + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w", encoding="utf-8", newline="") as file: + writer = csv.DictWriter(file, fieldnames=list(rows[0].keys())) + writer.writeheader() + writer.writerows(rows) + print(f"Wrote {output_path}") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/policyengine_us_data/datasets/cps/long_term/calibration.py b/policyengine_us_data/datasets/cps/long_term/calibration.py index e9d353727..f3c701cf9 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration.py @@ -1,5 +1,19 @@ import numpy as np import pandas as pd +from scipy import optimize, sparse + + +def _pct_error(achieved, target): + if target == 0: + return 0.0 if achieved == 0 else float("inf") + return (achieved - target) / target * 100 + + +def _relative_errors(achieved, target): + target = np.asarray(target, dtype=float) + achieved = np.asarray(achieved, dtype=float) + denominator = np.maximum(np.abs(target), 1e-10) + return (achieved - target) / denominator def iterative_proportional_fitting( @@ -47,7 +61,7 @@ def iterative_proportional_fitting( if verbose and (iter_num % 10 == 0 or rel_change < tol): predictions_new = X.T @ w - rel_errors = np.abs(predictions_new - y) / y + rel_errors = np.abs(_relative_errors(predictions_new, y)) max_rel_error = rel_errors.max() print( f"Iteration {iter_num:3d}: Max relative error = {max_rel_error:.6f}, Weight change = {rel_change:.6e}" @@ -66,8 +80,8 @@ def iterative_proportional_fitting( "iterations": iter_num + 1, "predictions_initial": predictions_initial, "predictions_new": predictions_final, - "relative_errors_initial": (predictions_initial - y) / y, - "relative_errors_new": (predictions_final - y) / y, + "relative_errors_initial": _relative_errors(predictions_initial, y), + "relative_errors_new": _relative_errors(predictions_final, y), "weight_ratio": w / w_initial, } @@ -166,6 +180,321 @@ def calibrate_greg( return w_new, 1 +def _build_constraint_dataframe_and_controls( + X, + y_target, + *, + ss_values=None, + ss_target=None, + payroll_values=None, + payroll_target=None, + h6_income_values=None, + h6_revenue_target=None, + oasdi_tob_values=None, + oasdi_tob_target=None, + hi_tob_values=None, + hi_tob_target=None, + n_ages=86, +): + controls = {} + age_cols = {f"age_{i}": X[:, i].astype(float) for i in range(n_ages)} + aux_df = pd.DataFrame(age_cols) + + for age_idx in range(n_ages): + controls[f"age_{age_idx}"] = float(y_target[age_idx]) + + if ss_values is not None and ss_target is not None: + aux_df["ss_total"] = np.asarray(ss_values, dtype=float) + controls["ss_total"] = float(ss_target) + + if payroll_values is not None and payroll_target is not None: + aux_df["payroll_total"] = np.asarray(payroll_values, dtype=float) + controls["payroll_total"] = float(payroll_target) + + if h6_income_values is not None and h6_revenue_target is not None: + aux_df["h6_revenue"] = np.asarray(h6_income_values, dtype=float) + controls["h6_revenue"] = float(h6_revenue_target) + + if oasdi_tob_values is not None and oasdi_tob_target is not None: + aux_df["oasdi_tob"] = np.asarray(oasdi_tob_values, dtype=float) + controls["oasdi_tob"] = float(oasdi_tob_target) + + if hi_tob_values is not None and hi_tob_target is not None: + aux_df["hi_tob"] = np.asarray(hi_tob_values, dtype=float) + controls["hi_tob"] = float(hi_tob_target) + + return aux_df, controls + + +def calibrate_entropy( + X, + y_target, + baseline_weights, + ss_values=None, + ss_target=None, + payroll_values=None, + payroll_target=None, + h6_income_values=None, + h6_revenue_target=None, + oasdi_tob_values=None, + oasdi_tob_target=None, + hi_tob_values=None, + hi_tob_target=None, + n_ages=86, + max_iters=500, + tol=1e-10, +): + """ + Positive calibration via entropy balancing. + + Finds strictly positive weights minimizing KL divergence from the baseline + weights while matching all requested calibration constraints. + """ + + aux_df, controls = _build_constraint_dataframe_and_controls( + X, + y_target, + ss_values=ss_values, + ss_target=ss_target, + payroll_values=payroll_values, + payroll_target=payroll_target, + h6_income_values=h6_income_values, + h6_revenue_target=h6_revenue_target, + oasdi_tob_values=oasdi_tob_values, + oasdi_tob_target=oasdi_tob_target, + hi_tob_values=hi_tob_values, + hi_tob_target=hi_tob_target, + n_ages=n_ages, + ) + + A = aux_df.to_numpy(dtype=float) + targets = np.array(list(controls.values()), dtype=float) + scales = np.maximum( + np.maximum(np.abs(targets), np.abs(A.T @ baseline_weights)), + 1.0, + ) + A_scaled = A / scales + targets_scaled = targets / scales + + baseline_weights = np.asarray(baseline_weights, dtype=float) + gram = A_scaled.T @ (baseline_weights[:, None] * A_scaled) + gram += np.eye(gram.shape[0]) * 1e-12 + beta0 = np.linalg.solve(gram, targets_scaled - (A_scaled.T @ baseline_weights)) + + def objective_gradient_hessian(beta): + eta = np.clip(A_scaled @ beta, -700, 700) + exp_eta = np.exp(eta) + weights = baseline_weights * exp_eta + objective = float(np.sum(weights) - targets_scaled @ beta) + gradient = A_scaled.T @ weights - targets_scaled + hessian = A_scaled.T @ (weights[:, None] * A_scaled) + return objective, gradient, hessian + + def solve_with_root(beta_start): + result = optimize.root( + lambda z: objective_gradient_hessian(z)[1], + beta_start, + jac=lambda z: objective_gradient_hessian(z)[2], + method="hybr", + options={"xtol": tol}, + ) + if not result.success: + return None + _, gradient, _ = objective_gradient_hessian(result.x) + max_error = float( + np.max(100 * np.abs(gradient) / np.maximum(np.abs(targets_scaled), 1e-12)) + ) + if max_error > tol * 100: + return None + return result.x, result.nfev + + def infeasibility_error(prefix): + feasibility = assess_nonnegative_feasibility(A, targets) + if feasibility["success"]: + return RuntimeError( + f"{prefix}. Nonnegative exact calibration appears infeasible under current support; " + f"best achievable max relative constraint error is " + f"{feasibility['best_case_max_pct_error']:.3f}%." + ) + return RuntimeError( + f"{prefix}. Nonnegative feasibility diagnostic could not certify a solution: " + f"{feasibility['message']}" + ) + + beta = beta0.copy() + iterations = 0 + final_max_error = float("inf") + + for iterations in range(1, max_iters + 1): + objective, gradient, hessian = objective_gradient_hessian(beta) + final_max_error = float( + np.max(100 * np.abs(gradient) / np.maximum(np.abs(targets_scaled), 1e-12)) + ) + if final_max_error <= tol * 100: + break + + hessian += np.eye(hessian.shape[0]) * 1e-12 + try: + delta = np.linalg.solve(hessian, gradient) + except np.linalg.LinAlgError: + delta = np.linalg.lstsq(hessian, gradient, rcond=None)[0] + + step = 1.0 + while step >= 1e-8: + candidate = beta - step * delta + candidate_objective, candidate_gradient, _ = objective_gradient_hessian(candidate) + candidate_max_error = float( + np.max( + 100 + * np.abs(candidate_gradient) + / np.maximum(np.abs(targets_scaled), 1e-12) + ) + ) + if np.isfinite(candidate_objective) and ( + candidate_objective <= objective + 1e-12 + or candidate_max_error < final_max_error + ): + beta = candidate + break + step /= 2.0 + + if step < 1e-8: + root_solution = solve_with_root(beta) + if root_solution is not None: + beta, root_iterations = root_solution + iterations += int(root_iterations) + break + raise infeasibility_error( + "Entropy calibration line search failed to find a descent step" + ) + else: + root_solution = solve_with_root(beta) + if root_solution is None: + raise infeasibility_error( + "Entropy calibration failed: " + f"max constraint error remained {final_max_error:.6f}% " + f"after {max_iters} iterations" + ) + beta, root_iterations = root_solution + iterations += int(root_iterations) + + eta = np.clip(A_scaled @ beta, -700, 700) + weights = baseline_weights * np.exp(eta) + return weights, iterations + + +def calibrate_entropy_approximate( + X, + y_target, + baseline_weights, + ss_values=None, + ss_target=None, + payroll_values=None, + payroll_target=None, + h6_income_values=None, + h6_revenue_target=None, + oasdi_tob_values=None, + oasdi_tob_target=None, + hi_tob_values=None, + hi_tob_target=None, + n_ages=86, +): + """ + Approximate nonnegative calibration via minimax relative-error LP. + + This is the robust fallback when exact positive entropy calibration is + infeasible under the current support. It returns the best nonnegative + weight vector available under the requested constraints. + """ + + aux_df, controls = _build_constraint_dataframe_and_controls( + X, + y_target, + ss_values=ss_values, + ss_target=ss_target, + payroll_values=payroll_values, + payroll_target=payroll_target, + h6_income_values=h6_income_values, + h6_revenue_target=h6_revenue_target, + oasdi_tob_values=oasdi_tob_values, + oasdi_tob_target=oasdi_tob_target, + hi_tob_values=hi_tob_values, + hi_tob_target=hi_tob_target, + n_ages=n_ages, + ) + + A = aux_df.to_numpy(dtype=float) + targets = np.array(list(controls.values()), dtype=float) + feasibility = assess_nonnegative_feasibility(A, targets, return_weights=True) + weights = feasibility.get("weights") + if not feasibility["success"] or weights is None: + raise RuntimeError( + "Approximate nonnegative calibration failed: " + f"{feasibility['message']}" + ) + + return np.asarray(weights, dtype=float), 1, feasibility + + +def assess_nonnegative_feasibility(A, targets, *, return_weights=False): + """ + Solve for the minimum uniform relative error achievable with nonnegative weights. + + Returns a dict with `success` and `best_case_max_pct_error`. + """ + A = np.asarray(A, dtype=float) + targets = np.asarray(targets, dtype=float) + if A.shape[1] == len(targets): + constraint_by_unit = A.T + elif A.shape[0] == len(targets): + constraint_by_unit = A + else: + raise ValueError( + "Constraint matrix shape does not match targets: " + f"{A.shape} vs {targets.shape}" + ) + + scales = np.maximum(np.abs(targets), 1.0) + A_rel = constraint_by_unit / scales[:, None] + b_rel = targets / scales + + constraint_matrix = sparse.csr_matrix(A_rel) + epsilon_column = sparse.csc_matrix(np.ones((constraint_matrix.shape[0], 1))) + A_ub = sparse.vstack( + [ + sparse.hstack([constraint_matrix, -epsilon_column]), + sparse.hstack([-constraint_matrix, -epsilon_column]), + ], + format="csc", + ) + b_ub = np.concatenate([b_rel, -b_rel]) + c = np.zeros(constraint_matrix.shape[1] + 1) + c[-1] = 1.0 + bounds = [(0, None)] * constraint_matrix.shape[1] + [(0, None)] + + result = optimize.linprog( + c, + A_ub=A_ub, + b_ub=b_ub, + bounds=bounds, + method="highs", + ) + + result_dict = { + "success": bool(result.success), + "best_case_max_pct_error": ( + float(result.x[-1] * 100) if result.success else None + ), + "status": int(result.status), + "message": result.message, + } + if return_weights: + result_dict["weights"] = ( + np.asarray(result.x[:-1], dtype=float) if result.success else None + ) + return result_dict + + def calibrate_weights( X, y_target, @@ -186,6 +515,9 @@ def calibrate_weights( max_iters=100, tol=1e-6, verbose=False, + allow_fallback_to_ipf=True, + allow_approximate_entropy=False, + approximate_max_error_pct=None, ): """ Unified interface for weight calibration. @@ -214,12 +546,25 @@ def calibrate_weights( Returns: w_new: Calibrated weights iterations: Number of iterations + audit: Metadata about calibration method selection """ + audit = { + "method_requested": method, + "method_used": method, + "greg_attempted": method == "greg", + "greg_error": None, + "fell_back_to_ipf": False, + "lp_fallback_used": False, + "approximate_solution_used": False, + "approximation_method": None, + "approximate_solution_error_pct": None, + } + if method == "greg": if calibrator is None: raise ValueError("calibrator required for GREG method") try: - return calibrate_greg( + w_new, iterations = calibrate_greg( calibrator, X, y_target, @@ -236,15 +581,154 @@ def calibrate_weights( hi_tob_target, n_ages, ) + return w_new, iterations, audit except Exception as e: + audit["greg_error"] = str(e) + if not allow_fallback_to_ipf: + raise RuntimeError( + "GREG calibration failed while fallback was disabled" + ) from e if verbose: print(f"GREG failed: {e}, falling back to IPF") w_new, info = iterative_proportional_fitting( X, y_target, baseline_weights, max_iters, tol, verbose ) - return w_new, info["iterations"] + audit["method_used"] = "ipf" + audit["fell_back_to_ipf"] = True + return w_new, info["iterations"], audit + elif method == "entropy": + try: + w_new, iterations = calibrate_entropy( + X, + y_target, + baseline_weights, + ss_values=ss_values, + ss_target=ss_target, + payroll_values=payroll_values, + payroll_target=payroll_target, + h6_income_values=h6_income_values, + h6_revenue_target=h6_revenue_target, + oasdi_tob_values=oasdi_tob_values, + oasdi_tob_target=oasdi_tob_target, + hi_tob_values=hi_tob_values, + hi_tob_target=hi_tob_target, + n_ages=n_ages, + max_iters=max_iters * 5, + tol=max(tol, 1e-10), + ) + return w_new, iterations, audit + except RuntimeError as error: + audit["entropy_error"] = str(error) + w_new, iterations, feasibility = calibrate_entropy_approximate( + X, + y_target, + baseline_weights, + ss_values=ss_values, + ss_target=ss_target, + payroll_values=payroll_values, + payroll_target=payroll_target, + h6_income_values=h6_income_values, + h6_revenue_target=h6_revenue_target, + oasdi_tob_values=oasdi_tob_values, + oasdi_tob_target=oasdi_tob_target, + hi_tob_values=hi_tob_values, + hi_tob_target=hi_tob_target, + n_ages=n_ages, + ) + approximate_error_pct = float(feasibility["best_case_max_pct_error"]) + if approximate_error_pct <= max(tol * 100, 1e-6): + audit["lp_fallback_used"] = True + audit["approximation_method"] = "lp_minimax_exact" + audit["approximate_solution_error_pct"] = approximate_error_pct + return w_new, iterations, audit + + if not allow_approximate_entropy: + raise + + if ( + approximate_max_error_pct is not None + and approximate_error_pct > approximate_max_error_pct + ): + raise RuntimeError( + "Approximate entropy fallback exceeded allowable error: " + f"{approximate_error_pct:.3f}% > {approximate_max_error_pct:.3f}%" + ) from error + + audit["lp_fallback_used"] = True + audit["approximate_solution_used"] = True + audit["approximation_method"] = "lp_minimax" + audit["approximate_solution_error_pct"] = approximate_error_pct + return w_new, iterations, audit else: w_new, info = iterative_proportional_fitting( X, y_target, baseline_weights, max_iters, tol, verbose ) - return w_new, info["iterations"] + return w_new, info["iterations"], audit + + +def build_calibration_audit( + *, + X, + y_target, + weights, + baseline_weights, + calibration_event, + ss_values=None, + ss_target=None, + payroll_values=None, + payroll_target=None, + h6_income_values=None, + h6_revenue_target=None, + oasdi_tob_values=None, + oasdi_tob_target=None, + hi_tob_values=None, + hi_tob_target=None, +): + achieved_ages = X.T @ weights + age_errors = np.abs(achieved_ages - y_target) / np.maximum(np.abs(y_target), 1e-10) * 100 + + neg_mask = weights < 0 + negative_values = np.abs(weights[neg_mask]) + + audit = dict(calibration_event) + audit.update( + { + "age_max_pct_error": float(age_errors.max()), + "negative_weight_count": int(neg_mask.sum()), + "negative_weight_pct": float(100 * neg_mask.sum() / len(weights)), + "largest_negative_weight": float(negative_values.max()) if negative_values.size else 0.0, + "constraints": {}, + "baseline_weight_sum": float(np.sum(baseline_weights)), + "calibrated_weight_sum": float(np.sum(weights)), + "max_constraint_pct_error": 0.0, + } + ) + + constraint_specs = [ + ("ss_total", ss_values, ss_target), + ("payroll_total", payroll_values, payroll_target), + ("h6_revenue", h6_income_values, h6_revenue_target), + ("oasdi_tob", oasdi_tob_values, oasdi_tob_target), + ("hi_tob", hi_tob_values, hi_tob_target), + ] + + for name, values, target in constraint_specs: + if values is None or target is None: + continue + achieved = float(np.sum(values * weights)) + audit["constraints"][name] = { + "target": float(target), + "achieved": achieved, + "error": achieved - float(target), + "pct_error": float(_pct_error(achieved, float(target))), + } + + if audit["constraints"]: + audit["max_constraint_pct_error"] = float( + max( + abs(stats["pct_error"]) + for stats in audit["constraints"].values() + ) + ) + + return audit diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py new file mode 100644 index 000000000..789b4d8ac --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +from datetime import datetime, timezone +import json +from pathlib import Path +from typing import Any + +try: + from .calibration_profiles import classify_calibration_quality, get_profile +except ImportError: # pragma: no cover - script execution fallback + from calibration_profiles import classify_calibration_quality, get_profile + + +CONTRACT_VERSION = 1 +MANIFEST_FILENAME = "calibration_manifest.json" + + +def metadata_path_for(h5_path: str | Path) -> Path: + return Path(f"{Path(h5_path)}.metadata.json") + + +def normalize_metadata(metadata: dict[str, Any]) -> dict[str, Any]: + normalized = json.loads(json.dumps(metadata)) + normalized.setdefault("contract_version", CONTRACT_VERSION) + + profile_data = normalized.get("profile", {}) + audit = normalized.setdefault("calibration_audit", {}) + constraints = audit.get("constraints", {}) + + if "max_constraint_pct_error" not in audit: + audit["max_constraint_pct_error"] = float( + max((abs(stats.get("pct_error", 0.0)) for stats in constraints.values()), default=0.0) + ) + + if "calibration_quality" not in audit and profile_data.get("name"): + try: + profile = get_profile(profile_data["name"]) + except ValueError: + profile = None + if profile is not None: + canonical_profile = profile.to_dict() + merged_profile = json.loads(json.dumps(canonical_profile)) + merged_profile.update(profile_data) + normalized["profile"] = merged_profile + audit["calibration_quality"] = classify_calibration_quality( + audit, + profile, + year=normalized.get("year"), + ) + + return normalized + + +def write_year_metadata( + h5_path: str | Path, + *, + year: int, + base_dataset_path: str, + profile: dict[str, Any], + calibration_audit: dict[str, Any], + target_source: dict[str, Any] | None = None, +) -> Path: + metadata = { + "contract_version": CONTRACT_VERSION, + "year": year, + "base_dataset_path": base_dataset_path, + "profile": profile, + "calibration_audit": calibration_audit, + } + if target_source is not None: + metadata["target_source"] = target_source + metadata = normalize_metadata(metadata) + metadata_path = metadata_path_for(h5_path) + metadata_path.write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return metadata_path + + +def update_dataset_manifest( + output_dir: str | Path, + *, + year: int, + h5_path: str | Path, + metadata_path: str | Path, + base_dataset_path: str, + profile: dict[str, Any], + calibration_audit: dict[str, Any], + target_source: dict[str, Any] | None = None, +) -> Path: + output_dir = Path(output_dir) + manifest_path = output_dir / MANIFEST_FILENAME + profile = json.loads(json.dumps(profile)) + target_source = json.loads(json.dumps(target_source)) + + if manifest_path.exists(): + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + else: + manifest = { + "contract_version": CONTRACT_VERSION, + "generated_at": None, + "base_dataset_path": base_dataset_path, + "profile": profile, + "target_source": target_source, + "years": [], + "datasets": {}, + } + + if manifest["base_dataset_path"] != base_dataset_path: + raise ValueError( + "Output directory already contains a different base dataset path: " + f"{manifest['base_dataset_path']} != {base_dataset_path}" + ) + manifest_profile = json.loads(json.dumps(manifest["profile"])) + if manifest_profile != profile: + if ( + manifest_profile.get("name") == profile.get("name") + and manifest_profile.get("calibration_method") + == profile.get("calibration_method") + ): + manifest["profile"] = profile + else: + raise ValueError( + "Output directory already contains a different calibration profile: " + f"{manifest['profile'].get('name')} != {profile.get('name')}" + ) + if manifest.get("target_source") is None and target_source is not None: + manifest["target_source"] = target_source + elif manifest.get("target_source") != target_source: + raise ValueError( + "Output directory already contains a different target source: " + f"{manifest.get('target_source')} != {target_source}" + ) + + datasets = manifest.setdefault("datasets", {}) + datasets[str(year)] = { + "h5": Path(h5_path).name, + "metadata": Path(metadata_path).name, + "calibration_quality": calibration_audit.get("calibration_quality"), + "method_used": calibration_audit.get("method_used"), + "fell_back_to_ipf": calibration_audit.get("fell_back_to_ipf"), + "age_max_pct_error": calibration_audit.get("age_max_pct_error"), + "max_constraint_pct_error": calibration_audit.get("max_constraint_pct_error"), + "negative_weight_pct": calibration_audit.get("negative_weight_pct"), + } + + year_set = {int(value) for value in manifest.get("years", [])} + year_set.add(year) + manifest["years"] = sorted(year_set) + manifest["year_range"] = { + "start": min(year_set), + "end": max(year_set), + } + manifest["generated_at"] = datetime.now(timezone.utc).isoformat() + + manifest_path.write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return manifest_path + + +def rebuild_dataset_manifest(output_dir: str | Path) -> Path: + return rebuild_dataset_manifest_with_target_source(output_dir) + + +def rebuild_dataset_manifest_with_target_source( + output_dir: str | Path, + *, + target_source: dict[str, Any] | None = None, +) -> Path: + output_dir = Path(output_dir) + metadata_files = sorted(output_dir.glob("*.h5.metadata.json")) + if not metadata_files: + raise FileNotFoundError(f"No metadata sidecars found in {output_dir}") + + manifest_path: Path | None = None + for metadata_file in metadata_files: + metadata = json.loads(metadata_file.read_text(encoding="utf-8")) + metadata = normalize_metadata(metadata) + if target_source is not None: + metadata["target_source"] = target_source + metadata_file.write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + year = int(metadata["year"]) + h5_path = output_dir / f"{year}.h5" + manifest_path = update_dataset_manifest( + output_dir, + year=year, + h5_path=h5_path, + metadata_path=metadata_file, + base_dataset_path=metadata["base_dataset_path"], + profile=metadata["profile"], + calibration_audit=metadata["calibration_audit"], + target_source=metadata.get("target_source"), + ) + + assert manifest_path is not None + return manifest_path diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py new file mode 100644 index 000000000..420916824 --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py @@ -0,0 +1,379 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from typing import Any + + +@dataclass(frozen=True) +class ApproximateCalibrationWindow: + start_year: int + end_year: int | None + max_constraint_error_pct: float + max_age_error_pct: float + max_negative_weight_pct: float | None = 0.0 + + def applies(self, year: int) -> bool: + if year < self.start_year: + return False + if self.end_year is not None and year > self.end_year: + return False + return True + + +@dataclass(frozen=True) +class CalibrationProfile: + name: str + description: str + calibration_method: str + use_greg: bool + use_ss: bool + use_payroll: bool + use_h6_reform: bool + use_tob: bool + allow_greg_fallback: bool = False + max_constraint_error_pct: float = 0.1 + max_age_error_pct: float = 0.1 + max_negative_weight_pct: float | None = None + approximate_windows: tuple[ApproximateCalibrationWindow, ...] = field( + default_factory=tuple + ) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +DEFAULT_LONG_RUN_APPROXIMATE_WINDOWS = ( + ApproximateCalibrationWindow( + start_year=2075, + end_year=2078, + max_constraint_error_pct=0.5, + max_age_error_pct=0.5, + max_negative_weight_pct=0.0, + ), + ApproximateCalibrationWindow( + start_year=2079, + end_year=2085, + max_constraint_error_pct=5.0, + max_age_error_pct=5.0, + max_negative_weight_pct=0.0, + ), + ApproximateCalibrationWindow( + start_year=2086, + end_year=2095, + max_constraint_error_pct=20.0, + max_age_error_pct=20.0, + max_negative_weight_pct=0.0, + ), + ApproximateCalibrationWindow( + start_year=2096, + end_year=None, + max_constraint_error_pct=35.0, + max_age_error_pct=35.0, + max_negative_weight_pct=0.0, + ), +) + + +NAMED_PROFILES: dict[str, CalibrationProfile] = { + "age-only": CalibrationProfile( + name="age-only", + description="Age-only calibration using IPF.", + calibration_method="ipf", + use_greg=False, + use_ss=False, + use_payroll=False, + use_h6_reform=False, + use_tob=False, + allow_greg_fallback=False, + ), + "ss": CalibrationProfile( + name="ss", + description="Age plus Social Security benefits using positive entropy calibration.", + calibration_method="entropy", + use_greg=False, + use_ss=True, + use_payroll=False, + use_h6_reform=False, + use_tob=False, + max_negative_weight_pct=0.0, + approximate_windows=DEFAULT_LONG_RUN_APPROXIMATE_WINDOWS, + ), + "ss-payroll": CalibrationProfile( + name="ss-payroll", + description="Age, Social Security, and taxable payroll using positive entropy calibration.", + calibration_method="entropy", + use_greg=False, + use_ss=True, + use_payroll=True, + use_h6_reform=False, + use_tob=False, + max_negative_weight_pct=0.0, + approximate_windows=DEFAULT_LONG_RUN_APPROXIMATE_WINDOWS, + ), + "ss-payroll-tob": CalibrationProfile( + name="ss-payroll-tob", + description="Age, Social Security, taxable payroll, and TOB using positive entropy calibration.", + calibration_method="entropy", + use_greg=False, + use_ss=True, + use_payroll=True, + use_h6_reform=False, + use_tob=True, + max_negative_weight_pct=0.0, + approximate_windows=DEFAULT_LONG_RUN_APPROXIMATE_WINDOWS, + ), + "ss-payroll-tob-h6": CalibrationProfile( + name="ss-payroll-tob-h6", + description="Age, Social Security, taxable payroll, TOB, and H6 using positive entropy calibration.", + calibration_method="entropy", + use_greg=False, + use_ss=True, + use_payroll=True, + use_h6_reform=True, + use_tob=True, + max_negative_weight_pct=0.0, + approximate_windows=DEFAULT_LONG_RUN_APPROXIMATE_WINDOWS, + ), +} + +QUALITY_RANK = { + "aggregate": 0, + "approximate": 1, + "exact": 2, +} + + +def get_profile(name: str) -> CalibrationProfile: + try: + return NAMED_PROFILES[name] + except KeyError as error: + valid = ", ".join(sorted(NAMED_PROFILES)) + raise ValueError(f"Unknown calibration profile '{name}'. Valid profiles: {valid}") from error + + +def approximate_window_for_year( + profile: CalibrationProfile, + year: int | None, +) -> ApproximateCalibrationWindow | None: + if not profile.approximate_windows: + return None + + if year is None: + return max( + profile.approximate_windows, + key=lambda window: ( + float("-inf") if window.end_year is None else window.end_year, + window.max_constraint_error_pct, + window.max_age_error_pct, + ), + ) + + for window in profile.approximate_windows: + if window.applies(year): + return window + return None + + +def build_profile_from_flags( + *, + use_greg: bool, + use_ss: bool, + use_payroll: bool, + use_h6_reform: bool, + use_tob: bool, +) -> CalibrationProfile: + if not use_greg: + for profile in NAMED_PROFILES.values(): + if ( + profile.use_greg is False + and profile.use_ss == use_ss + and profile.use_payroll == use_payroll + and profile.use_h6_reform == use_h6_reform + and profile.use_tob == use_tob + ): + return profile + + for profile in NAMED_PROFILES.values(): + if ( + profile.calibration_method == ("greg" if use_greg else "ipf") + and profile.use_greg == use_greg + and profile.use_ss == use_ss + and profile.use_payroll == use_payroll + and profile.use_h6_reform == use_h6_reform + and profile.use_tob == use_tob + ): + return profile + + flag_names = [] + if use_greg: + flag_names.append("greg") + if use_ss: + flag_names.append("ss") + if use_payroll: + flag_names.append("payroll") + if use_h6_reform: + flag_names.append("h6") + if use_tob: + flag_names.append("tob") + + suffix = "-".join(flag_names) if flag_names else "age-only" + return CalibrationProfile( + name=f"custom-{suffix}", + description="Legacy flag-derived calibration profile.", + calibration_method="greg" if use_greg else "ipf", + use_greg=use_greg, + use_ss=use_ss, + use_payroll=use_payroll, + use_h6_reform=use_h6_reform, + use_tob=use_tob, + ) + + +def validate_calibration_audit( + audit: dict[str, Any], + profile: CalibrationProfile, + *, + year: int | None = None, + quality: str | None = None, +) -> list[str]: + if quality is None: + quality = audit.get("calibration_quality") or classify_calibration_quality( + audit, + profile, + year=year, + ) + + if quality == "exact": + return _collect_threshold_issues( + audit, + profile, + max_constraint_error_pct=profile.max_constraint_error_pct, + max_age_error_pct=profile.max_age_error_pct, + max_negative_weight_pct=profile.max_negative_weight_pct, + ) + + if quality == "approximate": + window = approximate_window_for_year(profile, year) + if window is None: + issues = _collect_threshold_issues( + audit, + profile, + max_constraint_error_pct=profile.max_constraint_error_pct, + max_age_error_pct=profile.max_age_error_pct, + max_negative_weight_pct=profile.max_negative_weight_pct, + ) + issues.append( + "Approximate calibration is not permitted for this profile/year" + ) + return issues + return _collect_threshold_issues( + audit, + profile, + max_constraint_error_pct=window.max_constraint_error_pct, + max_age_error_pct=window.max_age_error_pct, + max_negative_weight_pct=window.max_negative_weight_pct, + ) + + exact_issues = _collect_threshold_issues( + audit, + profile, + max_constraint_error_pct=profile.max_constraint_error_pct, + max_age_error_pct=profile.max_age_error_pct, + max_negative_weight_pct=profile.max_negative_weight_pct, + ) + window = approximate_window_for_year(profile, year) + if window is None: + return exact_issues + [ + "Calibration quality aggregate exceeds approximate thresholds" + ] + approximate_issues = _collect_threshold_issues( + audit, + profile, + max_constraint_error_pct=window.max_constraint_error_pct, + max_age_error_pct=window.max_age_error_pct, + max_negative_weight_pct=window.max_negative_weight_pct, + ) + return approximate_issues + [ + "Calibration quality aggregate exceeds approximate thresholds" + ] + + +def _collect_threshold_issues( + audit: dict[str, Any], + profile: CalibrationProfile, + *, + max_constraint_error_pct: float | None, + max_age_error_pct: float | None, + max_negative_weight_pct: float | None, +) -> list[str]: + issues: list[str] = [] + + if profile.calibration_method == "greg" and audit.get("fell_back_to_ipf"): + issues.append("GREG calibration fell back to IPF") + + age_error = audit.get("age_max_pct_error") + if ( + max_age_error_pct is not None + and age_error is not None + and age_error > max_age_error_pct + ): + issues.append( + f"Age max error {age_error:.3f}% exceeds {max_age_error_pct:.3f}%" + ) + + for constraint_name, stats in audit.get("constraints", {}).items(): + pct_error = stats.get("pct_error") + if ( + max_constraint_error_pct is not None + and pct_error is not None + and abs(pct_error) > max_constraint_error_pct + ): + issues.append( + f"{constraint_name} error {pct_error:.3f}% exceeds " + f"{max_constraint_error_pct:.3f}%" + ) + + if max_negative_weight_pct is not None: + pct = audit.get("negative_weight_pct") + if pct is not None and pct > max_negative_weight_pct: + issues.append( + f"Negative weight share {pct:.3f}% exceeds " + f"{max_negative_weight_pct:.3f}%" + ) + + return issues + + +def classify_calibration_quality( + audit: dict[str, Any], + profile: CalibrationProfile, + *, + year: int | None = None, +) -> str: + exact_issues = _collect_threshold_issues( + audit, + profile, + max_constraint_error_pct=profile.max_constraint_error_pct, + max_age_error_pct=profile.max_age_error_pct, + max_negative_weight_pct=profile.max_negative_weight_pct, + ) + if not exact_issues: + return "exact" + + window = approximate_window_for_year(profile, year) + if window is None: + if year is not None: + return "aggregate" + return "approximate" + + approximate_issues = _collect_threshold_issues( + audit, + profile, + max_constraint_error_pct=window.max_constraint_error_pct, + max_age_error_pct=window.max_age_error_pct, + max_negative_weight_pct=window.max_negative_weight_pct, + ) + if not approximate_issues: + return "approximate" + + return "aggregate" diff --git a/policyengine_us_data/datasets/cps/long_term/rebuild_calibration_manifest.py b/policyengine_us_data/datasets/cps/long_term/rebuild_calibration_manifest.py new file mode 100644 index 000000000..bcae97125 --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/rebuild_calibration_manifest.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +import argparse + +from calibration_artifacts import rebuild_dataset_manifest_with_target_source +from ssa_data import describe_long_term_target_source + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Rebuild calibration_manifest.json from existing metadata sidecars.", + ) + parser.add_argument( + "output_dir", + help="Directory containing YYYY.h5 and YYYY.h5.metadata.json files.", + ) + parser.add_argument( + "--target-source", + help="Optional target source name to stamp into each sidecar while rebuilding the manifest.", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + target_source = ( + describe_long_term_target_source(args.target_source) + if args.target_source + else None + ) + manifest_path = rebuild_dataset_manifest_with_target_source( + args.output_dir, + target_source=target_source, + ) + print(f"Rebuilt {manifest_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py index 1413efe4b..9d7f78db0 100644 --- a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py +++ b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py @@ -3,10 +3,14 @@ Usage: + python run_household_projection.py [START_YEAR] [END_YEAR] [--profile PROFILE] [--target-source SOURCE] [--output-dir DIR] [--save-h5] python run_household_projection.py [START_YEAR] [END_YEAR] [--greg] [--use-ss] [--use-payroll] [--use-h6-reform] [--use-tob] [--save-h5] START_YEAR: Optional starting year (default: 2025) END_YEAR: Optional ending year (default: 2035) + --profile: Named calibration contract (recommended) + --target-source: Named long-term target source package + --output-dir: Output directory for generated H5 files and metadata --greg: Use GREG calibration instead of IPF (optional) --use-ss: Include Social Security benefit totals as calibration target (requires --greg) --use-payroll: Include taxable payroll totals as calibration target (requires --greg) @@ -15,8 +19,8 @@ --save-h5: Save year-specific .h5 files with calibrated weights to ./projected_datasets/ Examples: - python run_household_projection.py 2045 2045 --greg --use-ss # single year - python run_household_projection.py 2025 2100 --greg --use-ss --use-payroll --use-tob --save-h5 + python run_household_projection.py 2045 2045 --profile ss --target-source trustees_2025_current_law --save-h5 + python run_household_projection.py 2025 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --save-h5 """ import sys @@ -29,11 +33,22 @@ from policyengine_us import Microsimulation from ssa_data import ( + describe_long_term_target_source, + get_long_term_target_source, load_ssa_age_projections, load_ssa_benefit_projections, load_taxable_payroll_projections, + set_long_term_target_source, +) +from calibration import build_calibration_audit, calibrate_weights +from calibration_artifacts import update_dataset_manifest, write_year_metadata +from calibration_profiles import ( + approximate_window_for_year, + build_profile_from_flags, + classify_calibration_quality, + get_profile, + validate_calibration_audit, ) -from calibration import calibrate_weights from projection_utils import ( build_household_age_matrix, create_household_year_h5, @@ -227,6 +242,31 @@ def create_h6_reform(): BASE_YEAR = DATASET_OPTIONS[SELECTED_DATASET]["base_year"] +PROFILE_NAME = None +if "--profile" in sys.argv: + profile_index = sys.argv.index("--profile") + if profile_index + 1 >= len(sys.argv): + raise ValueError("--profile requires a profile name") + PROFILE_NAME = sys.argv[profile_index + 1] + del sys.argv[profile_index : profile_index + 2] + +TARGET_SOURCE = None +if "--target-source" in sys.argv: + source_index = sys.argv.index("--target-source") + if source_index + 1 >= len(sys.argv): + raise ValueError("--target-source requires a source name") + TARGET_SOURCE = sys.argv[source_index + 1] + del sys.argv[source_index : source_index + 2] + +OUTPUT_DIR = "./projected_datasets" +if "--output-dir" in sys.argv: + output_dir_index = sys.argv.index("--output-dir") + if output_dir_index + 1 >= len(sys.argv): + raise ValueError("--output-dir requires a directory path") + OUTPUT_DIR = sys.argv[output_dir_index + 1] + del sys.argv[output_dir_index : output_dir_index + 2] + + USE_GREG = "--greg" in sys.argv if USE_GREG: sys.argv.remove("--greg") @@ -251,7 +291,6 @@ def create_h6_reform(): if not USE_GREG: print("Warning: --use-h6-reform requires --greg, enabling GREG automatically") USE_GREG = True - from ssa_data import load_h6_income_rate_change USE_TOB = "--use-tob" in sys.argv if USE_TOB: @@ -259,7 +298,6 @@ def create_h6_reform(): if not USE_GREG: print("Warning: --use-tob requires --greg, enabling GREG automatically") USE_GREG = True - from ssa_data import load_oasdi_tob_projections, load_hi_tob_projections SAVE_H5 = "--save-h5" in sys.argv if SAVE_H5: @@ -268,6 +306,39 @@ def create_h6_reform(): START_YEAR = int(sys.argv[1]) if len(sys.argv) > 1 else 2025 END_YEAR = int(sys.argv[2]) if len(sys.argv) > 2 else 2035 +legacy_flags_used = any([USE_GREG, USE_SS, USE_PAYROLL, USE_H6_REFORM, USE_TOB]) +if PROFILE_NAME and legacy_flags_used: + raise ValueError("Use either --profile or legacy calibration flags, not both.") + +if PROFILE_NAME: + PROFILE = get_profile(PROFILE_NAME) +else: + PROFILE = build_profile_from_flags( + use_greg=USE_GREG, + use_ss=USE_SS, + use_payroll=USE_PAYROLL, + use_h6_reform=USE_H6_REFORM, + use_tob=USE_TOB, + ) + +if TARGET_SOURCE: + set_long_term_target_source(TARGET_SOURCE) +TARGET_SOURCE = get_long_term_target_source() +TARGET_SOURCE_METADATA = describe_long_term_target_source(TARGET_SOURCE) + +CALIBRATION_METHOD = PROFILE.calibration_method +USE_GREG = CALIBRATION_METHOD == "greg" +USE_SS = PROFILE.use_ss +USE_PAYROLL = PROFILE.use_payroll +USE_H6_REFORM = PROFILE.use_h6_reform +USE_TOB = PROFILE.use_tob + +if USE_H6_REFORM: + from ssa_data import load_h6_income_rate_change + +if USE_TOB: + from ssa_data import load_hi_tob_projections, load_oasdi_tob_projections + if USE_GREG: try: from samplics.weighting import SampleWeight @@ -280,8 +351,6 @@ def create_h6_reform(): else: calibrator = None -OUTPUT_DIR = "./projected_datasets" - print("=" * 70) print(f"HOUSEHOLD-LEVEL INCOME TAX PROJECTION: {START_YEAR}-{END_YEAR}") print("=" * 70) @@ -289,7 +358,10 @@ def create_h6_reform(): print(f" Base year: {BASE_YEAR} (CPS microdata)") print(f" Projection: {START_YEAR}-{END_YEAR}") print(f" Calculation level: HOUSEHOLD ONLY (simplified)") -print(f" Calibration method: {'GREG' if USE_GREG else 'IPF'}") +print(f" Calibration profile: {PROFILE.name}") +print(f" Profile description: {PROFILE.description}") +print(f" Target source: {TARGET_SOURCE}") +print(f" Calibration method: {CALIBRATION_METHOD.upper()}") if USE_SS: print(f" Including Social Security benefits constraint: Yes") if USE_PAYROLL: @@ -497,12 +569,13 @@ def create_h6_reform(): ) y_target = target_matrix[:, year_idx] + approximate_window = approximate_window_for_year(PROFILE, year) - w_new, iterations = calibrate_weights( + w_new, iterations, calibration_event = calibrate_weights( X=X, y_target=y_target, baseline_weights=baseline_weights, - method="greg" if USE_GREG else "ipf", + method=CALIBRATION_METHOD, calibrator=calibrator, ss_values=ss_values, ss_target=ss_target, @@ -518,14 +591,52 @@ def create_h6_reform(): max_iters=100, tol=1e-6, verbose=False, + allow_fallback_to_ipf=PROFILE.allow_greg_fallback, + allow_approximate_entropy=approximate_window is not None, + approximate_max_error_pct=( + approximate_window.max_constraint_error_pct + if approximate_window is not None + else None + ), ) - if year in display_years and USE_GREG: - neg_mask = w_new < 0 - n_neg = neg_mask.sum() + calibration_audit = build_calibration_audit( + X=X, + y_target=y_target, + weights=w_new, + baseline_weights=baseline_weights, + calibration_event=calibration_event, + ss_values=ss_values, + ss_target=ss_target, + payroll_values=payroll_values, + payroll_target=payroll_target, + h6_income_values=h6_income_values, + h6_revenue_target=h6_revenue_target, + oasdi_tob_values=oasdi_tob_values, + oasdi_tob_target=oasdi_tob_target, + hi_tob_values=hi_tob_values, + hi_tob_target=hi_tob_target, + ) + calibration_audit["calibration_quality"] = classify_calibration_quality( + calibration_audit, + PROFILE, + year=year, + ) + + validation_issues = validate_calibration_audit( + calibration_audit, + PROFILE, + year=year, + ) + if validation_issues: + issue_text = "; ".join(validation_issues) + raise RuntimeError(f"Calibration validation failed for {year}: {issue_text}") + + if year in display_years and CALIBRATION_METHOD in {"greg", "entropy"}: + n_neg = calibration_audit["negative_weight_count"] if n_neg > 0: - pct_neg = 100 * n_neg / len(w_new) - max_neg = np.abs(w_new[neg_mask]).max() + pct_neg = calibration_audit["negative_weight_pct"] + max_neg = calibration_audit["largest_negative_weight"] print( f" [DEBUG {year}] Negative weights: {n_neg} ({pct_neg:.2f}%), " f"largest: {max_neg:,.0f}" @@ -535,33 +646,38 @@ def create_h6_reform(): if year in display_years and (USE_SS or USE_PAYROLL or USE_H6_REFORM or USE_TOB): if USE_SS: - ss_achieved = np.sum(ss_values * w_new) + ss_stats = calibration_audit["constraints"]["ss_total"] print( - f" [DEBUG {year}] SS achieved: ${ss_achieved / 1e9:.1f}B (error: ${abs(ss_achieved - ss_target) / 1e6:.1f}M, {(ss_achieved - ss_target) / ss_target * 100:.3f}%)" + f" [DEBUG {year}] SS achieved: ${ss_stats['achieved'] / 1e9:.1f}B " + f"(error: ${abs(ss_stats['error']) / 1e6:.1f}M, " + f"{ss_stats['pct_error']:.3f}%)" ) if USE_PAYROLL: - payroll_achieved = np.sum(payroll_values * w_new) + payroll_stats = calibration_audit["constraints"]["payroll_total"] print( - f" [DEBUG {year}] Payroll achieved: ${payroll_achieved / 1e9:.1f}B (error: ${abs(payroll_achieved - payroll_target) / 1e6:.1f}M, {(payroll_achieved - payroll_target) / payroll_target * 100:.3f}%)" + f" [DEBUG {year}] Payroll achieved: ${payroll_stats['achieved'] / 1e9:.1f}B " + f"(error: ${abs(payroll_stats['error']) / 1e6:.1f}M, " + f"{payroll_stats['pct_error']:.3f}%)" ) if USE_H6_REFORM and h6_revenue_target is not None: - h6_revenue_achieved = np.sum(h6_income_values * w_new) - error_pct = ( - (h6_revenue_achieved - h6_revenue_target) / abs(h6_revenue_target) * 100 - if h6_revenue_target != 0 - else 0 - ) + h6_stats = calibration_audit["constraints"]["h6_revenue"] print( - f" [DEBUG {year}] H6 achieved revenue: ${h6_revenue_achieved / 1e9:.3f}B (error: ${abs(h6_revenue_achieved - h6_revenue_target) / 1e6:.1f}M, {error_pct:.3f}%)" + f" [DEBUG {year}] H6 achieved revenue: ${h6_stats['achieved'] / 1e9:.3f}B " + f"(error: ${abs(h6_stats['error']) / 1e6:.1f}M, " + f"{h6_stats['pct_error']:.3f}%)" ) if USE_TOB: - oasdi_achieved = np.sum(oasdi_tob_values * w_new) - hi_achieved = np.sum(hi_tob_values * w_new) + oasdi_stats = calibration_audit["constraints"]["oasdi_tob"] + hi_stats = calibration_audit["constraints"]["hi_tob"] print( - f" [DEBUG {year}] OASDI TOB achieved: ${oasdi_achieved / 1e9:.1f}B (error: ${abs(oasdi_achieved - oasdi_tob_target) / 1e6:.1f}M, {(oasdi_achieved - oasdi_tob_target) / oasdi_tob_target * 100:.3f}%)" + f" [DEBUG {year}] OASDI TOB achieved: ${oasdi_stats['achieved'] / 1e9:.1f}B " + f"(error: ${abs(oasdi_stats['error']) / 1e6:.1f}M, " + f"{oasdi_stats['pct_error']:.3f}%)" ) print( - f" [DEBUG {year}] HI TOB achieved: ${hi_achieved / 1e9:.1f}B (error: ${abs(hi_achieved - hi_tob_target) / 1e6:.1f}M, {(hi_achieved - hi_tob_target) / hi_tob_target * 100:.3f}%)" + f" [DEBUG {year}] HI TOB achieved: ${hi_stats['achieved'] / 1e9:.1f}B " + f"(error: ${abs(hi_stats['error']) / 1e6:.1f}M, " + f"{hi_stats['pct_error']:.3f}%)" ) weights_matrix[:, year_idx] = w_new @@ -572,8 +688,26 @@ def create_h6_reform(): if SAVE_H5: h5_path = create_household_year_h5(year, w_new, BASE_DATASET_PATH, OUTPUT_DIR) + metadata_path = write_year_metadata( + h5_path, + year=year, + base_dataset_path=BASE_DATASET_PATH, + profile=PROFILE.to_dict(), + calibration_audit=calibration_audit, + target_source=TARGET_SOURCE_METADATA, + ) + update_dataset_manifest( + OUTPUT_DIR, + year=year, + h5_path=h5_path, + metadata_path=metadata_path, + base_dataset_path=BASE_DATASET_PATH, + profile=PROFILE.to_dict(), + calibration_audit=calibration_audit, + target_source=TARGET_SOURCE_METADATA, + ) if year in display_years: - print(f" Saved {year}.h5") + print(f" Saved {year}.h5 and metadata") del sim gc.collect() diff --git a/policyengine_us_data/datasets/cps/long_term/ssa_data.py b/policyengine_us_data/datasets/cps/long_term/ssa_data.py index 6b76ae219..b54ba0550 100644 --- a/policyengine_us_data/datasets/cps/long_term/ssa_data.py +++ b/policyengine_us_data/datasets/cps/long_term/ssa_data.py @@ -1,8 +1,77 @@ +import json +import os +from functools import lru_cache + import numpy as np import pandas as pd from policyengine_us_data.storage import STORAGE_FOLDER +LONG_TERM_TARGET_SOURCES_DIR = STORAGE_FOLDER / "long_term_target_sources" +LONG_TERM_TARGET_SOURCES_MANIFEST = LONG_TERM_TARGET_SOURCES_DIR / "sources.json" +DEFAULT_LONG_TERM_TARGET_SOURCE = "trustees_2025_current_law" +_CURRENT_LONG_TERM_TARGET_SOURCE = os.environ.get( + "POLICYENGINE_US_DATA_LONG_TERM_TARGET_SOURCE", + DEFAULT_LONG_TERM_TARGET_SOURCE, +) + + +@lru_cache(maxsize=1) +def _load_long_term_target_sources_manifest() -> dict: + return json.loads(LONG_TERM_TARGET_SOURCES_MANIFEST.read_text(encoding="utf-8")) + + +def available_long_term_target_sources() -> list[str]: + manifest = _load_long_term_target_sources_manifest() + return sorted(manifest["sources"]) + + +def get_long_term_target_source() -> str: + return _CURRENT_LONG_TERM_TARGET_SOURCE + + +def set_long_term_target_source(source_name: str) -> None: + global _CURRENT_LONG_TERM_TARGET_SOURCE + _CURRENT_LONG_TERM_TARGET_SOURCE = resolve_long_term_target_source_name(source_name) + + +def resolve_long_term_target_source_name(source_name: str | None = None) -> str: + manifest = _load_long_term_target_sources_manifest() + candidate = source_name or _CURRENT_LONG_TERM_TARGET_SOURCE + if candidate not in manifest["sources"]: + valid = ", ".join(sorted(manifest["sources"])) + raise ValueError( + f"Unknown long-term target source {candidate!r}. Valid sources: {valid}" + ) + return candidate + + +def describe_long_term_target_source(source_name: str | None = None) -> dict: + manifest = _load_long_term_target_sources_manifest() + resolved_name = resolve_long_term_target_source_name(source_name) + source = dict(manifest["sources"][resolved_name]) + source["name"] = resolved_name + return source + + +@lru_cache(maxsize=None) +def _load_long_term_target_frame(source_name: str) -> pd.DataFrame: + source = describe_long_term_target_source(source_name) + csv_path = LONG_TERM_TARGET_SOURCES_DIR / source["file"] + return pd.read_csv(csv_path) + + +def _load_long_term_target_row(year: int, source_name: str | None = None) -> pd.Series: + resolved_name = resolve_long_term_target_source_name(source_name) + df = _load_long_term_target_frame(resolved_name) + row = df[df["year"] == year] + if row.empty: + raise ValueError( + f"Year {year} not found in long-term target source {resolved_name!r}" + ) + return row.iloc[0] + + def load_ssa_age_projections(start_year=2025, end_year=2100): """ Load SSA population projections from package storage. @@ -37,7 +106,7 @@ def load_ssa_age_projections(start_year=2025, end_year=2100): return target_matrix -def load_ssa_benefit_projections(year): +def load_ssa_benefit_projections(year, source_name: str | None = None): """ Load SSA Trustee Report projections for Social Security benefits. @@ -47,15 +116,12 @@ def load_ssa_benefit_projections(year): Returns: Total OASDI benefits in nominal dollars """ - csv_path = STORAGE_FOLDER / "social_security_aux.csv" - df = pd.read_csv(csv_path) - - row = df[df["year"] == year] - nominal_billions = row["oasdi_cost_in_billion_nominal_usd"].values[0] + row = _load_long_term_target_row(year, source_name) + nominal_billions = row["oasdi_cost_in_billion_nominal_usd"] return nominal_billions * 1e9 -def load_taxable_payroll_projections(year): +def load_taxable_payroll_projections(year, source_name: str | None = None): """ Load SSA Trustee Report projections for taxable payroll. @@ -65,15 +131,12 @@ def load_taxable_payroll_projections(year): Returns: Total taxable payroll in nominal dollars """ - csv_path = STORAGE_FOLDER / "social_security_aux.csv" - df = pd.read_csv(csv_path) - - row = df[df["year"] == year] - nominal_billions = row["taxable_payroll_in_billion_nominal_usd"].values[0] + row = _load_long_term_target_row(year, source_name) + nominal_billions = row["taxable_payroll_in_billion_nominal_usd"] return nominal_billions * 1e9 -def load_h6_income_rate_change(year): +def load_h6_income_rate_change(year, source_name: str | None = None): """ Load H6 reform income rate change target for a given year. @@ -83,15 +146,12 @@ def load_h6_income_rate_change(year): Returns: H6 income rate change as decimal (e.g., -0.0018 for -0.18%) """ - csv_path = STORAGE_FOLDER / "social_security_aux.csv" - df = pd.read_csv(csv_path) - - row = df[df["year"] == year] + row = _load_long_term_target_row(year, source_name) # CSV stores as percentage (e.g., -0.18), convert to decimal - return row["h6_income_rate_change"].values[0] / 100 + return row["h6_income_rate_change"] / 100 -def load_oasdi_tob_projections(year): +def load_oasdi_tob_projections(year, source_name: str | None = None): """ Load OASDI TOB (Taxation of Benefits) revenue target for a given year. @@ -101,15 +161,12 @@ def load_oasdi_tob_projections(year): Returns: Total OASDI TOB revenue in nominal dollars """ - csv_path = STORAGE_FOLDER / "social_security_aux.csv" - df = pd.read_csv(csv_path) - - row = df[df["year"] == year] - nominal_billions = row["oasdi_tob_billions_nominal_usd"].values[0] + row = _load_long_term_target_row(year, source_name) + nominal_billions = row["oasdi_tob_billions_nominal_usd"] return nominal_billions * 1e9 -def load_hi_tob_projections(year): +def load_hi_tob_projections(year, source_name: str | None = None): """ Load HI (Medicare) TOB revenue target for a given year. @@ -119,9 +176,6 @@ def load_hi_tob_projections(year): Returns: Total HI TOB revenue in nominal dollars """ - csv_path = STORAGE_FOLDER / "social_security_aux.csv" - df = pd.read_csv(csv_path) - - row = df[df["year"] == year] - nominal_billions = row["hi_tob_billions_nominal_usd"].values[0] + row = _load_long_term_target_row(year, source_name) + nominal_billions = row["hi_tob_billions_nominal_usd"] return nominal_billions * 1e9 diff --git a/policyengine_us_data/storage/README.md b/policyengine_us_data/storage/README.md index 80d4c1cdc..ac4478902 100644 --- a/policyengine_us_data/storage/README.md +++ b/policyengine_us_data/storage/README.md @@ -30,6 +30,13 @@ - `https://www.ssa.gov/oact/solvency/provisions/tables/table_run133.html` • Notes: Contains OASDI cost projections and taxable payroll data (2025-2100) +- **long_term_target_sources/** + • Source packages for long-term CPS calibration targets + • Files: + - `trustees_2025_current_law.csv`: explicit frozen copy of the legacy Trustees/current-law target path + - `sources.json`: provenance and source metadata for each named package + • Notes: `run_household_projection.py --target-source ...` selects from these packages instead of relying on branch-specific data files + - **national_and_district_rents_2023.csv** • Source: Census ACS 5-year estimates (2023), median 2BR rent by congressional district • Created by: `fetch_cd_rents.py` (requires `CENSUS_API_KEY` environment variable) diff --git a/policyengine_us_data/storage/long_term_target_sources/sources.json b/policyengine_us_data/storage/long_term_target_sources/sources.json new file mode 100644 index 000000000..3b1932697 --- /dev/null +++ b/policyengine_us_data/storage/long_term_target_sources/sources.json @@ -0,0 +1,18 @@ +{ + "default_source": "trustees_2025_current_law", + "sources": { + "trustees_2025_current_law": { + "description": "2025 Trustees current-law baseline used by the legacy long-term calibration stack.", + "file": "trustees_2025_current_law.csv", + "name": "trustees_2025_current_law", + "notes": [ + "Generated from social_security_aux.csv for explicit source selection." + ], + "source_urls": [ + "https://www.ssa.gov/oact/tr/2025/lrIndex.html", + "https://www.ssa.gov/oact/solvency/provisions/tables/table_run133.html" + ], + "type": "trustees_current_law" + } + } +} diff --git a/policyengine_us_data/storage/long_term_target_sources/trustees_2025_current_law.csv b/policyengine_us_data/storage/long_term_target_sources/trustees_2025_current_law.csv new file mode 100644 index 000000000..cc5d66108 --- /dev/null +++ b/policyengine_us_data/storage/long_term_target_sources/trustees_2025_current_law.csv @@ -0,0 +1,77 @@ +year,oasdi_cost_in_billion_2025_usd,cpi_w_intermediate,oasdi_cost_in_billion_nominal_usd,taxable_payroll_in_billion_nominal_usd,h6_income_rate_change,oasdi_tob_pct_of_taxable_payroll,oasdi_tob_billions_nominal_usd,hi_tob_billions_nominal_usd +2025,1609,100.0,1609.0,10621,0.0,0.57,60.5397,40.655 +2026,1660,102.49,1701.334,11129,0.0,0.69,76.7901,52.199 +2027,1715,104.95,1799.8925,11627,0.0,0.71,82.5517,60.569 +2028,1763,107.47,1894.6961,12159,0.0,0.73,88.7607,65.635 +2029,1810,110.05,1991.905,12696,0.0,0.75,95.22,70.797 +2030,1856,112.69,2091.5264,13239,0.0,0.78,103.2642,76.594 +2031,1903,115.4,2196.062,13798,0.0,0.81,111.7638,83.036 +2032,1947,118.17,2300.7699,14380,0.0,0.84,120.792,89.882 +2033,1991,121.0,2409.11,14987,0.0,0.86,128.8882,97.047 +2034,2032,123.91,2517.8512,15594,0.0,0.89,138.7866,104.521 +2035,2073,126.88,2630.2224,16205,0.0,0.9,145.845,111.683 +2036,2114,129.93,2746.7202,16825,0.0,0.91,153.1075,118.34 +2037,2155,133.04,2867.012,17465,0.0,0.91,158.9315,124.952 +2038,2194,136.24,2989.1056,18132,0.0,0.92,166.8144,131.665 +2039,2233,139.51,3115.2583,18819,0.0,0.93,175.0167,138.444 +2040,2270,142.86,3242.922,19532,0.0,0.94,183.6008,145.274 +2041,2306,146.28,3373.2168,20269,0.0,0.94,190.5286,152.175 +2042,2342,149.79,3508.0818,21035,0.0,0.94,197.729,159.209 +2043,2378,153.39,3647.6142,21828,0.0,0.95,207.366,166.412 +2044,2415,157.07,3793.2405,22653,0.0,0.95,215.2035,173.854 +2045,2452,160.84,3943.7968,23507,-0.07,0.96,225.6672,181.587 +2046,2488,164.7,4097.736,24391,-0.12,0.96,234.1536,189.555 +2047,2527,168.65,4261.7855,25313,-0.18,0.96,243.0048,197.851 +2048,2567,172.7,4433.209,26270,-0.23,0.97,254.819,206.586 +2049,2609,176.85,4614.0165,27263,-0.27,0.97,264.4511,215.751 +2050,2652,181.09,4802.5068,28300,-0.32,0.98,277.34,230.026 +2051,2696,185.44,4999.4624,29376,-0.36,0.98,287.8848,240.281 +2052,2743,189.89,5208.6827,30494,-0.4,0.99,301.8906,251.078 +2053,2792,194.44,5428.7648,31661,-0.43,0.99,313.4439,262.512 +2054,2842,199.11,5658.7062,32869,-1.0,1.0,328.69,274.469 +2055,2895,203.89,5902.6155,34124,-1.01,1.01,344.6524,287.161 +2056,2950,208.78,6159.01,35432,-1.01,1.01,357.8632,300.557 +2057,3007,213.79,6428.6653,36790,-1.02,1.02,375.258,314.676 +2058,3066,218.93,6712.3938,38201,-1.03,1.03,393.4703,329.538 +2059,3125,224.18,7005.625,39670,-1.04,1.04,412.568,344.937 +2060,3184,229.56,7309.1904,41196,-1.04,1.04,428.4384,360.893 +2061,3243,235.07,7623.3201,42782,-1.05,1.05,449.211,377.293 +2062,3303,240.71,7950.6513,44429,-1.06,1.06,470.9474,394.399 +2063,3362,246.49,8286.9938,46136,-1.06,1.06,489.0416,412.059 +2064,3422,252.4,8637.128,47902,-1.07,1.07,512.5514,430.272 +2065,3483,258.46,9002.1618,49733,-1.07,1.07,532.1431,449.287 +2066,3544,264.66,9379.5504,51631,-1.08,1.08,557.6148,469.05 +2067,3607,271.02,9775.6914,53598,-1.09,1.09,584.2182,489.656 +2068,3670,277.52,10184.984,55637,-1.09,1.09,606.4433,511.168 +2069,3735,284.18,10614.123,57746,-1.1,1.1,635.206,533.596 +2070,3801,291.0,11060.91,59930,-1.1,1.1,659.23,556.96 +2071,3867,297.99,11523.2733,62196,-1.11,1.11,690.3756,581.15 +2072,3934,305.14,12004.2076,64543,-1.12,1.12,722.8816,606.266 +2073,4002,312.46,12504.6492,66975,-1.12,1.12,750.12,632.641 +2074,4071,319.96,13025.5716,69501,-1.13,1.13,785.3613,659.863 +2075,4139,327.64,13561.0196,72131,-1.13,1.13,815.0803,688.165 +2076,4206,335.5,14111.13,74862,-1.14,1.14,853.4268,717.049 +2077,4273,343.55,14679.8915,77698,-1.14,1.14,885.7572,746.777 +2078,4339,351.8,15264.602,80650,-1.14,1.14,919.41,777.205 +2079,4403,360.24,15861.3672,83727,-1.15,1.15,962.8605,808.543 +2080,4467,368.89,16478.3163,86933,-1.15,1.15,999.7295,840.797 +2081,4530,377.74,17111.622,90268,-1.15,1.15,1038.082,873.88 +2082,4593,386.81,17766.1833,93749,-1.15,1.15,1078.1135,907.958 +2083,4655,396.09,18437.9895,97381,-1.15,1.15,1119.8815,943.051 +2084,4716,405.6,19128.096,101163,-1.15,1.15,1163.3745,978.965 +2085,4775,415.33,19832.0075,105104,-1.15,1.15,1208.696,1015.209 +2086,4833,425.3,20554.749,109217,-1.14,1.14,1245.0738,1052.674 +2087,4891,435.51,21300.7941,113504,-1.14,1.14,1293.9456,1090.789 +2088,4948,445.96,22066.1008,117973,-1.14,1.14,1344.8922,1130.016 +2089,5006,456.66,22860.3996,122629,-1.13,1.13,1385.7077,1170.571 +2090,5064,467.62,23680.2768,127477,-1.13,1.13,1440.4901,1212.549 +2091,5125,478.84,24540.55,132518,-1.13,1.13,1497.4534,1256.274 +2092,5188,490.34,25438.8392,137764,-1.12,1.12,1542.9568,1302.101 +2093,5254,502.1,26380.334,143215,-1.12,1.12,1604.008,1350.009 +2094,5323,514.16,27368.7368,148876,-1.12,1.12,1667.4112,1400.378 +2095,5396,526.49,28409.4004,154754,-1.12,1.12,1733.2448,1453.666 +2096,5472,539.13,29501.1936,160855,-1.12,1.12,1801.576,1509.442 +2097,5551,552.07,30645.4057,167185,-1.11,1.11,1855.7535,1567.878 +2098,5633,565.32,31844.4756,173750,-1.11,1.11,1928.625,1629.435 +2099,5719,578.89,33106.7191,180557,-1.12,1.12,2022.2384,1694.187 +2100,5809,592.78,34434.5902,187614,-1.12,1.12,2101.2768,1761.512 diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py new file mode 100644 index 000000000..e83d5bf0b --- /dev/null +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -0,0 +1,390 @@ +from __future__ import annotations + +import json +import numpy as np +import pytest + +from policyengine_us_data.datasets.cps.long_term import calibration as calibration_module +from policyengine_us_data.datasets.cps.long_term.calibration import ( + assess_nonnegative_feasibility, + build_calibration_audit, + calibrate_entropy, + calibrate_weights, +) +from policyengine_us_data.datasets.cps.long_term.calibration_artifacts import ( + rebuild_dataset_manifest, + update_dataset_manifest, + write_year_metadata, +) +from policyengine_us_data.datasets.cps.long_term.calibration_profiles import ( + build_profile_from_flags, + classify_calibration_quality, + get_profile, + validate_calibration_audit, +) +from policyengine_us_data.datasets.cps.long_term.ssa_data import ( + available_long_term_target_sources, + describe_long_term_target_source, + load_taxable_payroll_projections, +) + + +class ExplodingCalibrator: + def calibrate(self, **kwargs): + raise RuntimeError("boom") + + +def test_named_profile_lookup(): + profile = get_profile("ss-payroll-tob") + assert profile.calibration_method == "entropy" + assert profile.use_greg is False + assert profile.use_ss is True + assert profile.use_payroll is True + assert profile.use_tob is True + assert profile.use_h6_reform is False + assert profile.max_negative_weight_pct == 0.0 + + +def test_legacy_flags_map_to_named_profile(): + profile = build_profile_from_flags( + use_greg=False, + use_ss=True, + use_payroll=True, + use_h6_reform=False, + use_tob=True, + ) + assert profile.name == "ss-payroll-tob" + assert profile.calibration_method == "entropy" + + +def test_strict_greg_failure_raises(): + X = np.array([[1.0, 0.0], [0.0, 1.0]]) + y_target = np.array([1.0, 1.0]) + baseline_weights = np.array([1.0, 1.0]) + + with pytest.raises(RuntimeError, match="fallback was disabled"): + calibrate_weights( + X=X, + y_target=y_target, + baseline_weights=baseline_weights, + method="greg", + calibrator=ExplodingCalibrator(), + allow_fallback_to_ipf=False, + ) + + +def test_build_calibration_audit_reports_constraint_error(): + X = np.array([[1.0, 0.0], [0.0, 1.0]]) + y_target = np.array([1.0, 1.0]) + baseline_weights = np.array([1.0, 1.0]) + weights = np.array([1.0, 1.0]) + audit = build_calibration_audit( + X=X, + y_target=y_target, + weights=weights, + baseline_weights=baseline_weights, + calibration_event={ + "method_requested": "greg", + "method_used": "greg", + "greg_attempted": True, + "greg_error": None, + "fell_back_to_ipf": False, + }, + payroll_values=np.array([10.0, 0.0]), + payroll_target=20.0, + ) + + assert audit["constraints"]["payroll_total"]["achieved"] == 10.0 + assert audit["constraints"]["payroll_total"]["pct_error"] == -50.0 + + +def test_profile_validation_rejects_fallback_and_large_error(): + profile = build_profile_from_flags( + use_greg=True, + use_ss=True, + use_payroll=True, + use_h6_reform=False, + use_tob=False, + ) + audit = { + "fell_back_to_ipf": True, + "age_max_pct_error": 0.0, + "negative_weight_pct": 0.0, + "constraints": { + "payroll_total": {"pct_error": 0.2}, + }, + } + + issues = validate_calibration_audit(audit, profile) + assert "GREG calibration fell back to IPF" in issues + assert any("payroll_total error" in issue for issue in issues) + + +def test_classify_calibration_quality_marks_invalid_audit_approximate(): + profile = get_profile("ss-payroll-tob") + quality = classify_calibration_quality( + { + "fell_back_to_ipf": False, + "age_max_pct_error": 0.0, + "negative_weight_pct": 0.0, + "constraints": { + "ss_total": {"pct_error": 0.0}, + "payroll_total": {"pct_error": 0.0}, + "oasdi_tob": {"pct_error": 0.5}, + "hi_tob": {"pct_error": 0.0}, + }, + }, + profile, + year=2078, + ) + assert quality == "approximate" + + +def test_entropy_profile_rejects_negative_weights(): + profile = get_profile("ss-payroll-tob") + issues = validate_calibration_audit( + { + "fell_back_to_ipf": False, + "age_max_pct_error": 0.0, + "negative_weight_pct": 0.01, + "constraints": { + "ss_total": {"pct_error": 0.0}, + "payroll_total": {"pct_error": 0.0}, + "oasdi_tob": {"pct_error": 0.0}, + "hi_tob": {"pct_error": 0.0}, + }, + }, + profile, + ) + assert any("Negative weight share" in issue for issue in issues) + + +def test_approximate_window_is_year_bounded(): + profile = get_profile("ss-payroll-tob") + quality = classify_calibration_quality( + { + "fell_back_to_ipf": False, + "age_max_pct_error": 3.0, + "negative_weight_pct": 0.0, + "constraints": { + "ss_total": {"pct_error": 0.0}, + "payroll_total": {"pct_error": 3.0}, + "oasdi_tob": {"pct_error": 3.0}, + "hi_tob": {"pct_error": 3.0}, + }, + }, + profile, + year=2080, + ) + assert quality == "approximate" + + quality = classify_calibration_quality( + { + "fell_back_to_ipf": False, + "age_max_pct_error": 3.0, + "negative_weight_pct": 0.0, + "constraints": { + "ss_total": {"pct_error": 0.0}, + "payroll_total": {"pct_error": 3.0}, + "oasdi_tob": {"pct_error": 3.0}, + "hi_tob": {"pct_error": 3.0}, + }, + }, + profile, + year=2035, + ) + assert quality == "aggregate" + + +def test_manifest_updates_and_rejects_profile_mismatch(tmp_path): + profile = get_profile("ss-payroll-tob") + audit = { + "method_used": "greg", + "fell_back_to_ipf": False, + "negative_weight_pct": 1.5, + } + + year_2026 = tmp_path / "2026.h5" + year_2026.write_text("", encoding="utf-8") + metadata_2026 = write_year_metadata( + year_2026, + year=2026, + base_dataset_path="hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5", + profile=profile.to_dict(), + calibration_audit=audit, + ) + manifest_path = update_dataset_manifest( + tmp_path, + year=2026, + h5_path=year_2026, + metadata_path=metadata_2026, + base_dataset_path="hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5", + profile=profile.to_dict(), + calibration_audit=audit, + ) + + year_2027 = tmp_path / "2027.h5" + year_2027.write_text("", encoding="utf-8") + metadata_2027 = write_year_metadata( + year_2027, + year=2027, + base_dataset_path="hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5", + profile=profile.to_dict(), + calibration_audit=audit, + ) + update_dataset_manifest( + tmp_path, + year=2027, + h5_path=year_2027, + metadata_path=metadata_2027, + base_dataset_path="hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5", + profile=profile.to_dict(), + calibration_audit=audit, + ) + + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + assert manifest["profile"]["name"] == "ss-payroll-tob" + assert manifest["years"] == [2026, 2027] + assert manifest["datasets"]["2026"]["metadata"] == "2026.h5.metadata.json" + + with pytest.raises(ValueError, match="different calibration profile"): + update_dataset_manifest( + tmp_path, + year=2028, + h5_path=tmp_path / "2028.h5", + metadata_path=tmp_path / "2028.h5.metadata.json", + base_dataset_path="hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5", + profile=get_profile("ss").to_dict(), + calibration_audit=audit, + ) + + manifest_path.unlink() + rebuilt_path = rebuild_dataset_manifest(tmp_path) + rebuilt = json.loads(rebuilt_path.read_text(encoding="utf-8")) + assert rebuilt["years"] == [2026, 2027] + + +def test_entropy_calibration_produces_nonnegative_weights_and_hits_targets(): + X = np.array( + [ + [1.0, 0.0], + [0.0, 1.0], + [1.0, 1.0], + ] + ) + y_target = np.array([2.0, 3.0]) + baseline_weights = np.array([1.0, 1.0, 1.0]) + payroll_values = np.array([1.0, 0.0, 2.0]) + payroll_target = 3.5 + + weights, _ = calibrate_entropy( + X=X, + y_target=y_target, + baseline_weights=baseline_weights, + payroll_values=payroll_values, + payroll_target=payroll_target, + n_ages=2, + ) + + assert np.all(weights > 0) + np.testing.assert_allclose(X.T @ weights, y_target, rtol=1e-8, atol=1e-8) + np.testing.assert_allclose(np.dot(payroll_values, weights), payroll_target, rtol=1e-8, atol=1e-8) + + +def test_entropy_calibration_can_fall_back_to_lp_approximate_solution(): + X = np.array( + [ + [1.0, 0.0, 1.0], + [0.0, 1.0, 1.0], + ] + ) + y_target = np.array([1.0, 1.0, 3.0]) + baseline_weights = np.array([1.0, 1.0]) + + weights, _, audit = calibrate_weights( + X=X, + y_target=y_target, + baseline_weights=baseline_weights, + method="entropy", + n_ages=3, + allow_approximate_entropy=True, + approximate_max_error_pct=40.0, + ) + + assert audit["approximate_solution_used"] is True + assert audit["approximation_method"] == "lp_minimax" + assert audit["approximate_solution_error_pct"] > 10.0 + assert np.all(weights >= 0) + + +def test_entropy_calibration_uses_lp_exact_fallback_even_before_approximate_window( + monkeypatch, +): + monkeypatch.setattr( + calibration_module, + "calibrate_entropy", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("entropy stalled")), + ) + monkeypatch.setattr( + calibration_module, + "calibrate_entropy_approximate", + lambda *args, **kwargs: ( + np.array([1.0, 2.0]), + 1, + {"best_case_max_pct_error": 0.0}, + ), + ) + + weights, _, audit = calibrate_weights( + X=np.array([[1.0], [0.0]]), + y_target=np.array([1.0]), + baseline_weights=np.array([1.0, 1.0]), + method="entropy", + n_ages=1, + allow_approximate_entropy=False, + ) + + np.testing.assert_allclose(weights, np.array([1.0, 2.0])) + assert audit["lp_fallback_used"] is True + assert audit["approximate_solution_used"] is False + assert audit["approximation_method"] == "lp_minimax_exact" + + +def test_nonnegative_feasibility_diagnostic_distinguishes_feasible_and_infeasible(): + feasible_A = np.array( + [ + [1.0, 0.0], + [0.0, 1.0], + [1.0, 1.0], + ] + ) + feasible_targets = np.array([1.0, 2.0, 3.0]) + feasible = assess_nonnegative_feasibility(feasible_A, feasible_targets) + assert feasible["success"] is True + assert feasible["best_case_max_pct_error"] < 1e-6 + + infeasible_A = np.array( + [ + [1.0, 0.0], + [0.0, 1.0], + [1.0, 1.0], + ] + ) + infeasible_targets = np.array([1.0, 1.0, 3.0]) + infeasible = assess_nonnegative_feasibility(infeasible_A, infeasible_targets) + assert infeasible["success"] is True + assert infeasible["best_case_max_pct_error"] > 10.0 + + +def test_long_term_target_sources_are_available_and_distinct(): + sources = available_long_term_target_sources() + assert "trustees_2025_current_law" in sources + + trustees = describe_long_term_target_source("trustees_2025_current_law") + assert trustees["file"] == "trustees_2025_current_law.csv" + + payroll_2026 = load_taxable_payroll_projections( + 2026, + source_name="trustees_2025_current_law", + ) + assert payroll_2026 == pytest.approx(11_129_000_000_000.0) From 6bc34e0273e5ae5490d59d6f5d833b978ad2a2ed Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 31 Mar 2026 07:37:35 -0400 Subject: [PATCH 06/31] Add support diagnostics to long-run calibration audit --- .../datasets/cps/long_term/calibration.py | 21 ++++++++++- .../cps/long_term/calibration_artifacts.py | 15 ++++++++ .../cps/long_term/calibration_profiles.py | 4 +-- .../test_long_term_calibration_contract.py | 36 +++++++++++++++++++ 4 files changed, 73 insertions(+), 3 deletions(-) diff --git a/policyengine_us_data/datasets/cps/long_term/calibration.py b/policyengine_us_data/datasets/cps/long_term/calibration.py index f3c701cf9..7dded31cd 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration.py @@ -689,6 +689,20 @@ def build_calibration_audit( neg_mask = weights < 0 negative_values = np.abs(weights[neg_mask]) + positive_mask = weights > 0 + weight_sum = float(np.sum(weights)) + if weight_sum > 0: + sorted_weights = np.sort(weights) + top_10_weight_share_pct = float(sorted_weights[-10:].sum() / weight_sum * 100) + top_100_weight_share_pct = float(sorted_weights[-100:].sum() / weight_sum * 100) + else: + top_10_weight_share_pct = 0.0 + top_100_weight_share_pct = 0.0 + + if weight_sum > 0 and float(np.dot(weights, weights)) > 0: + effective_sample_size = float(weight_sum**2 / np.dot(weights, weights)) + else: + effective_sample_size = 0.0 audit = dict(calibration_event) audit.update( @@ -697,9 +711,14 @@ def build_calibration_audit( "negative_weight_count": int(neg_mask.sum()), "negative_weight_pct": float(100 * neg_mask.sum() / len(weights)), "largest_negative_weight": float(negative_values.max()) if negative_values.size else 0.0, + "positive_weight_count": int(positive_mask.sum()), + "positive_weight_pct": float(100 * positive_mask.sum() / len(weights)), + "effective_sample_size": effective_sample_size, + "top_10_weight_share_pct": top_10_weight_share_pct, + "top_100_weight_share_pct": top_100_weight_share_pct, "constraints": {}, "baseline_weight_sum": float(np.sum(baseline_weights)), - "calibrated_weight_sum": float(np.sum(weights)), + "calibrated_weight_sum": weight_sum, "max_constraint_pct_error": 0.0, } ) diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py index 789b4d8ac..d1b71c3fc 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py @@ -32,6 +32,12 @@ def normalize_metadata(metadata: dict[str, Any]) -> dict[str, Any]: max((abs(stats.get("pct_error", 0.0)) for stats in constraints.values()), default=0.0) ) + if audit.get("lp_fallback_used"): + realized_error = float(audit.get("max_constraint_pct_error", 0.0)) + stored_error = audit.get("approximate_solution_error_pct") + if stored_error is None or float(stored_error) < realized_error: + audit["approximate_solution_error_pct"] = realized_error + if "calibration_quality" not in audit and profile_data.get("name"): try: profile = get_profile(profile_data["name"]) @@ -48,6 +54,15 @@ def normalize_metadata(metadata: dict[str, Any]) -> dict[str, Any]: year=normalized.get("year"), ) + if audit.get("lp_fallback_used"): + quality = audit.get("calibration_quality") + if quality == "exact": + audit["approximation_method"] = "lp_minimax_exact" + audit["approximate_solution_used"] = False + elif quality == "approximate": + audit["approximation_method"] = "lp_minimax" + audit["approximate_solution_used"] = True + return normalized diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py index 420916824..f6b6f97d9 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py @@ -53,8 +53,8 @@ def to_dict(self) -> dict[str, Any]: ApproximateCalibrationWindow( start_year=2079, end_year=2085, - max_constraint_error_pct=5.0, - max_age_error_pct=5.0, + max_constraint_error_pct=10.0, + max_age_error_pct=10.0, max_negative_weight_pct=0.0, ), ApproximateCalibrationWindow( diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index e83d5bf0b..1f0699137 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -12,6 +12,7 @@ calibrate_weights, ) from policyengine_us_data.datasets.cps.long_term.calibration_artifacts import ( + normalize_metadata, rebuild_dataset_manifest, update_dataset_manifest, write_year_metadata, @@ -96,6 +97,11 @@ def test_build_calibration_audit_reports_constraint_error(): assert audit["constraints"]["payroll_total"]["achieved"] == 10.0 assert audit["constraints"]["payroll_total"]["pct_error"] == -50.0 + assert audit["positive_weight_count"] == 2 + assert audit["positive_weight_pct"] == 100.0 + assert audit["effective_sample_size"] == pytest.approx(2.0) + assert audit["top_10_weight_share_pct"] == pytest.approx(100.0) + assert audit["top_100_weight_share_pct"] == pytest.approx(100.0) def test_profile_validation_rejects_fallback_and_large_error(): @@ -196,6 +202,36 @@ def test_approximate_window_is_year_bounded(): assert quality == "aggregate" +def test_normalize_metadata_harmonizes_lp_fallback_labels(): + profile = get_profile("ss-payroll-tob") + metadata = normalize_metadata( + { + "year": 2075, + "profile": profile.to_dict(), + "calibration_audit": { + "lp_fallback_used": True, + "approximation_method": "lp_minimax_exact", + "approximate_solution_error_pct": 0.0, + "max_constraint_pct_error": 0.368, + "age_max_pct_error": 0.0, + "negative_weight_pct": 0.0, + "constraints": { + "ss_total": {"pct_error": 0.368}, + "payroll_total": {"pct_error": 0.0}, + "oasdi_tob": {"pct_error": 0.0}, + "hi_tob": {"pct_error": 0.0}, + }, + }, + } + ) + + audit = metadata["calibration_audit"] + assert audit["calibration_quality"] == "approximate" + assert audit["approximation_method"] == "lp_minimax" + assert audit["approximate_solution_used"] is True + assert audit["approximate_solution_error_pct"] == pytest.approx(0.368) + + def test_manifest_updates_and_rejects_profile_mismatch(tmp_path): profile = get_profile("ss-payroll-tob") audit = { From 4dfa539776d7efc0fa5cab35215ad65c3280d8f1 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 31 Mar 2026 08:14:30 -0400 Subject: [PATCH 07/31] Add late-year age aggregation for calibration --- .../cps/long_term/calibration_profiles.py | 5 ++ .../cps/long_term/projection_utils.py | 49 +++++++++++++++++++ .../cps/long_term/run_household_projection.py | 31 ++++++++++-- .../test_long_term_calibration_contract.py | 28 +++++++++++ 4 files changed, 109 insertions(+), 4 deletions(-) diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py index f6b6f97d9..49d13a612 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py @@ -11,6 +11,7 @@ class ApproximateCalibrationWindow: max_constraint_error_pct: float max_age_error_pct: float max_negative_weight_pct: float | None = 0.0 + age_bucket_size: int | None = None def applies(self, year: int) -> bool: if year < self.start_year: @@ -49,6 +50,7 @@ def to_dict(self) -> dict[str, Any]: max_constraint_error_pct=0.5, max_age_error_pct=0.5, max_negative_weight_pct=0.0, + age_bucket_size=5, ), ApproximateCalibrationWindow( start_year=2079, @@ -56,6 +58,7 @@ def to_dict(self) -> dict[str, Any]: max_constraint_error_pct=10.0, max_age_error_pct=10.0, max_negative_weight_pct=0.0, + age_bucket_size=5, ), ApproximateCalibrationWindow( start_year=2086, @@ -63,6 +66,7 @@ def to_dict(self) -> dict[str, Any]: max_constraint_error_pct=20.0, max_age_error_pct=20.0, max_negative_weight_pct=0.0, + age_bucket_size=5, ), ApproximateCalibrationWindow( start_year=2096, @@ -70,6 +74,7 @@ def to_dict(self) -> dict[str, Any]: max_constraint_error_pct=35.0, max_age_error_pct=35.0, max_negative_weight_pct=0.0, + age_bucket_size=5, ), ) diff --git a/policyengine_us_data/datasets/cps/long_term/projection_utils.py b/policyengine_us_data/datasets/cps/long_term/projection_utils.py index 8aee4f3b7..eb9e9aa44 100644 --- a/policyengine_us_data/datasets/cps/long_term/projection_utils.py +++ b/policyengine_us_data/datasets/cps/long_term/projection_utils.py @@ -39,6 +39,55 @@ def build_household_age_matrix(sim, n_ages=86): return X, household_ids_unique, hh_id_to_idx +def build_age_bins(n_ages=86, bucket_size=None): + """ + Build age-bucket ranges over the single-year age target vector. + + The final bucket always preserves the open-ended 85+ slot. + """ + if bucket_size is None or bucket_size <= 1: + return [(age_idx, age_idx + 1) for age_idx in range(n_ages)] + + bins = [] + upper_single_age = max(n_ages - 1, 0) + for start in range(0, upper_single_age, bucket_size): + end = min(start + bucket_size, upper_single_age) + bins.append((start, end)) + bins.append((upper_single_age, n_ages)) + return bins + + +def aggregate_household_age_matrix(X, age_bins): + """ + Aggregate a single-year household age matrix into coarser age buckets. + """ + if len(age_bins) == X.shape[1] and all(end - start == 1 for start, end in age_bins): + return X + return np.column_stack([X[:, start:end].sum(axis=1) for start, end in age_bins]) + + +def aggregate_age_targets(targets, age_bins): + """ + Aggregate age targets over the first axis. + + Accepts either a single target vector `(n_ages,)` or a matrix + `(n_ages, n_years)`. + """ + targets = np.asarray(targets, dtype=float) + if targets.ndim == 1: + return np.array( + [targets[start:end].sum() for start, end in age_bins], + dtype=float, + ) + + return np.vstack( + [ + targets[start:end, :].sum(axis=0) + for start, end in age_bins + ] + ) + + def get_pseudo_input_variables(sim): """ Identify variables that appear as inputs but aggregate calculated values. diff --git a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py index 9d7f78db0..8983a947b 100644 --- a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py +++ b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py @@ -50,6 +50,9 @@ validate_calibration_audit, ) from projection_utils import ( + aggregate_age_targets, + aggregate_household_age_matrix, + build_age_bins, build_household_age_matrix, create_household_year_h5, ) @@ -417,6 +420,7 @@ def create_h6_reform(): sim = Microsimulation(dataset=BASE_DATASET_PATH) X, household_ids_unique, hh_id_to_idx = build_household_age_matrix(sim, n_ages) n_households = len(household_ids_unique) +aggregated_age_cache: dict[int, tuple[np.ndarray, np.ndarray]] = {} print(f"\nLoaded {n_households:,} households") print(f"Household age matrix shape: {X.shape}") @@ -568,11 +572,28 @@ def create_h6_reform(): f" [DEBUG {year}] HI TOB baseline: ${hi_baseline / 1e9:.1f}B, target: ${hi_tob_target / 1e9:.1f}B" ) - y_target = target_matrix[:, year_idx] approximate_window = approximate_window_for_year(PROFILE, year) + age_bucket_size = ( + approximate_window.age_bucket_size + if approximate_window is not None + else None + ) + if age_bucket_size and age_bucket_size > 1: + if age_bucket_size not in aggregated_age_cache: + age_bins = build_age_bins(n_ages=n_ages, bucket_size=age_bucket_size) + aggregated_age_cache[age_bucket_size] = ( + aggregate_household_age_matrix(X, age_bins), + aggregate_age_targets(target_matrix, age_bins), + ) + X_current, aggregated_target_matrix = aggregated_age_cache[age_bucket_size] + y_target = aggregated_target_matrix[:, year_idx] + else: + X_current = X + y_target = target_matrix[:, year_idx] + age_bucket_size = 1 w_new, iterations, calibration_event = calibrate_weights( - X=X, + X=X_current, y_target=y_target, baseline_weights=baseline_weights, method=CALIBRATION_METHOD, @@ -587,7 +608,7 @@ def create_h6_reform(): oasdi_tob_target=oasdi_tob_target, hi_tob_values=hi_tob_values, hi_tob_target=hi_tob_target, - n_ages=n_ages, + n_ages=X_current.shape[1], max_iters=100, tol=1e-6, verbose=False, @@ -601,7 +622,7 @@ def create_h6_reform(): ) calibration_audit = build_calibration_audit( - X=X, + X=X_current, y_target=y_target, weights=w_new, baseline_weights=baseline_weights, @@ -622,6 +643,8 @@ def create_h6_reform(): PROFILE, year=year, ) + calibration_audit["age_bucket_size"] = age_bucket_size + calibration_audit["age_bucket_count"] = int(X_current.shape[1]) validation_issues = validate_calibration_audit( calibration_audit, diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index 1f0699137..8b23561c7 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -23,6 +23,11 @@ get_profile, validate_calibration_audit, ) +from policyengine_us_data.datasets.cps.long_term.projection_utils import ( + aggregate_age_targets, + aggregate_household_age_matrix, + build_age_bins, +) from policyengine_us_data.datasets.cps.long_term.ssa_data import ( available_long_term_target_sources, describe_long_term_target_source, @@ -44,6 +49,29 @@ def test_named_profile_lookup(): assert profile.use_tob is True assert profile.use_h6_reform is False assert profile.max_negative_weight_pct == 0.0 + assert profile.approximate_windows[0].age_bucket_size == 5 + + +def test_age_bin_helpers_preserve_population_totals(): + bins = build_age_bins(n_ages=86, bucket_size=5) + assert bins[0] == (0, 5) + assert bins[-1] == (85, 86) + + X = np.eye(86) + y = np.arange(86, dtype=float) + X_coarse = aggregate_household_age_matrix(X, bins) + y_coarse = aggregate_age_targets(y, bins) + + assert X_coarse.shape == (86, 18) + assert y_coarse.shape == (18,) + assert X_coarse.sum() == pytest.approx(X.sum()) + assert y_coarse.sum() == pytest.approx(y.sum()) + + target_matrix = np.column_stack([y, y * 2]) + aggregated_target_matrix = aggregate_age_targets(target_matrix, bins) + assert aggregated_target_matrix.shape == (18, 2) + assert aggregated_target_matrix[:, 0].sum() == pytest.approx(y.sum()) + assert aggregated_target_matrix[:, 1].sum() == pytest.approx((y * 2).sum()) def test_legacy_flags_map_to_named_profile(): From 2172962b35f34f15b57ca577c1d0482a9d9d0d66 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 31 Mar 2026 08:21:57 -0400 Subject: [PATCH 08/31] Add long-run calibration comparison tools --- .../profile_support_concentration.py | 150 +++++++++++++++ .../long_term/summarize_calibration_runs.py | 179 ++++++++++++++++++ 2 files changed, 329 insertions(+) create mode 100644 policyengine_us_data/datasets/cps/long_term/profile_support_concentration.py create mode 100644 policyengine_us_data/datasets/cps/long_term/summarize_calibration_runs.py diff --git a/policyengine_us_data/datasets/cps/long_term/profile_support_concentration.py b/policyengine_us_data/datasets/cps/long_term/profile_support_concentration.py new file mode 100644 index 000000000..b6102dad5 --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/profile_support_concentration.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import h5py +import numpy as np + + +SS_COMPONENTS = ( + "social_security_retirement", + "social_security_disability", + "social_security_survivors", + "social_security_dependents", +) +PAYROLL_COMPONENTS = ( + "employment_income_before_lsr", + "self_employment_income_before_lsr", +) + + +def _read_year_array(store: h5py.File, name: str, year: int) -> np.ndarray: + return store[name][str(year)][()] + + +def _build_household_lookup(household_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + order = np.argsort(household_ids) + sorted_ids = household_ids[order] + return sorted_ids, order + + +def _household_index(sorted_household_ids: np.ndarray, order: np.ndarray, person_household_ids: np.ndarray) -> np.ndarray: + positions = np.searchsorted(sorted_household_ids, person_household_ids) + if np.any(positions >= len(sorted_household_ids)): + raise ValueError("Person household ids exceed household id support") + matched = sorted_household_ids[positions] + if not np.array_equal(matched, person_household_ids): + raise ValueError("Person household ids do not match household-level ids") + return order[positions] + + +def _load_year(path: Path, year: int) -> dict[str, np.ndarray]: + with h5py.File(path, "r") as store: + household_ids = _read_year_array(store, "household_id", year).astype(np.int64) + household_weights = _read_year_array(store, "household_weight", year).astype(float) + person_household_ids = _read_year_array(store, "person_household_id", year).astype(np.int64) + ages = _read_year_array(store, "age", year).astype(float) + payroll = np.zeros_like(ages, dtype=float) + for component in PAYROLL_COMPONENTS: + payroll += _read_year_array(store, component, year).astype(float) + social_security = np.zeros_like(ages, dtype=float) + for component in SS_COMPONENTS: + social_security += _read_year_array(store, component, year).astype(float) + + sorted_ids, order = _build_household_lookup(household_ids) + person_household_index = _household_index(sorted_ids, order, person_household_ids) + + return { + "household_ids": household_ids, + "household_weights": household_weights, + "person_household_ids": person_household_ids, + "person_household_index": person_household_index, + "ages": ages, + "payroll": payroll, + "social_security": social_security, + } + + +def _effective_sample_size(weights: np.ndarray) -> float: + total = float(weights.sum()) + denom = float(np.dot(weights, weights)) + if total <= 0 or denom <= 0: + return 0.0 + return total**2 / denom + + +def _top_households(data: dict[str, np.ndarray], top_n: int) -> list[dict[str, object]]: + weights = data["household_weights"] + top_idx = np.argsort(weights)[-top_n:][::-1] + records: list[dict[str, object]] = [] + for idx in top_idx: + if weights[idx] <= 0: + continue + mask = data["person_household_index"] == idx + ages = np.sort(data["ages"][mask]).astype(int).tolist() + payroll_total = float(data["payroll"][mask].sum()) + social_security_total = float(data["social_security"][mask].sum()) + records.append( + { + "household_id": int(data["household_ids"][idx]), + "weight": float(weights[idx]), + "weight_share_pct": float(weights[idx] / weights.sum() * 100), + "ages": ages, + "payroll_proxy": payroll_total, + "social_security_total": social_security_total, + } + ) + return records + + +def profile_support(path: Path, year: int, *, top_n: int) -> dict[str, object]: + data = _load_year(path, year) + household_weights = data["household_weights"] + positive_mask = household_weights > 0 + sorted_weights = np.sort(household_weights) + person_weights = household_weights[data["person_household_index"]] + ages = data["ages"] + payroll = data["payroll"] + + overall_nonworking = payroll <= 0 + age_85_plus = ages >= 85 + + return { + "path": str(path), + "year": year, + "positive_household_count": int(positive_mask.sum()), + "positive_household_pct": float(positive_mask.mean() * 100), + "effective_sample_size": _effective_sample_size(household_weights), + "top_10_weight_share_pct": float(sorted_weights[-10:].sum() / household_weights.sum() * 100), + "top_100_weight_share_pct": float(sorted_weights[-100:].sum() / household_weights.sum() * 100), + "weighted_nonworking_share_pct": float( + person_weights[overall_nonworking].sum() / person_weights.sum() * 100 + ), + "weighted_nonworking_share_85_plus_pct": float( + person_weights[age_85_plus & overall_nonworking].sum() + / person_weights[age_85_plus].sum() + * 100 + ) + if person_weights[age_85_plus].sum() > 0 + else 0.0, + "top_households": _top_households(data, top_n), + } + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Profile late-year support concentration in projected household datasets." + ) + parser.add_argument("dataset", type=Path, help="Projected year-specific H5 dataset.") + parser.add_argument("year", type=int, help="Projection year stored in the dataset.") + parser.add_argument("--top-n", type=int, default=20, help="Number of top households to emit.") + args = parser.parse_args() + + report = profile_support(args.dataset, args.year, top_n=args.top_n) + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/policyengine_us_data/datasets/cps/long_term/summarize_calibration_runs.py b/policyengine_us_data/datasets/cps/long_term/summarize_calibration_runs.py new file mode 100644 index 000000000..da9535a4e --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/summarize_calibration_runs.py @@ -0,0 +1,179 @@ +from __future__ import annotations + +import argparse +import csv +import json +from pathlib import Path +import sys +from typing import Any + +from calibration_artifacts import normalize_metadata +from profile_support_concentration import profile_support + + +SUPPORT_FIELDS = ( + "positive_household_count", + "positive_household_pct", + "effective_sample_size", + "top_10_weight_share_pct", + "top_100_weight_share_pct", + "weighted_nonworking_share_pct", + "weighted_nonworking_share_85_plus_pct", +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Summarize and compare calibration quality across projected dataset directories." + ) + ) + parser.add_argument("left", type=Path, help="First projected dataset directory.") + parser.add_argument("right", type=Path, nargs="?", help="Optional second directory to compare.") + parser.add_argument( + "--years", + help="Optional comma-separated list of years to include. Defaults to all years found.", + ) + parser.add_argument( + "--profile-support", + action="store_true", + help="Compute support-concentration metrics from the H5 files when metadata is missing them.", + ) + parser.add_argument( + "--output", + type=Path, + help="Optional CSV output path. Prints to stdout when omitted.", + ) + return parser.parse_args() + + +def parse_years(raw: str | None) -> list[int] | None: + if not raw: + return None + return [int(value.strip()) for value in raw.split(",") if value.strip()] + + +def metadata_for(directory: Path, year: int) -> dict[str, Any] | None: + metadata_path = directory / f"{year}.h5.metadata.json" + if not metadata_path.exists(): + return None + return normalize_metadata(json.loads(metadata_path.read_text(encoding="utf-8"))) + + +def support_metrics(directory: Path, year: int, metadata: dict[str, Any] | None, *, profile: bool) -> dict[str, Any]: + audit = (metadata or {}).get("calibration_audit", {}) + metrics = { + field: audit.get(field) + for field in SUPPORT_FIELDS + } + if not profile or all(value is not None for value in metrics.values()): + return metrics + + dataset_path = directory / f"{year}.h5" + if not dataset_path.exists(): + return metrics + profiled = profile_support(dataset_path, year, top_n=10) + return { + "positive_household_count": profiled["positive_household_count"], + "positive_household_pct": profiled["positive_household_pct"], + "effective_sample_size": profiled["effective_sample_size"], + "top_10_weight_share_pct": profiled["top_10_weight_share_pct"], + "top_100_weight_share_pct": profiled["top_100_weight_share_pct"], + "weighted_nonworking_share_pct": profiled["weighted_nonworking_share_pct"], + "weighted_nonworking_share_85_plus_pct": profiled["weighted_nonworking_share_85_plus_pct"], + } + + +def summarize_directory(directory: Path, years: list[int] | None, *, profile: bool) -> dict[int, dict[str, Any]]: + if years is None: + years = sorted( + int(path.name.split(".")[0]) + for path in directory.glob("*.h5.metadata.json") + ) + + rows: dict[int, dict[str, Any]] = {} + for year in years: + metadata = metadata_for(directory, year) + if metadata is None: + continue + audit = metadata["calibration_audit"] + row: dict[str, Any] = { + "quality": audit.get("calibration_quality"), + "method": audit.get("approximation_method") or audit.get("method_used"), + "age_bucket_size": audit.get("age_bucket_size"), + "max_constraint_pct_error": audit.get("max_constraint_pct_error"), + "age_max_pct_error": audit.get("age_max_pct_error"), + "negative_weight_pct": audit.get("negative_weight_pct"), + } + row.update(support_metrics(directory, year, metadata, profile=profile)) + rows[year] = row + return rows + + +def build_rows( + left: dict[int, dict[str, Any]], + right: dict[int, dict[str, Any]] | None, +) -> list[dict[str, Any]]: + years = sorted(set(left) | set(right or {})) + rows: list[dict[str, Any]] = [] + for year in years: + row: dict[str, Any] = {"year": year} + for prefix, source in [("left", left), ("right", right or {})]: + values = source.get(year, {}) + for key, value in values.items(): + row[f"{prefix}_{key}"] = value + if right is not None and year in left and year in right: + for key in ( + "max_constraint_pct_error", + "age_max_pct_error", + "effective_sample_size", + "top_10_weight_share_pct", + "top_100_weight_share_pct", + "weighted_nonworking_share_pct", + "weighted_nonworking_share_85_plus_pct", + ): + left_value = left[year].get(key) + right_value = right[year].get(key) + if left_value is not None and right_value is not None: + row[f"delta_{key}"] = right_value - left_value + rows.append(row) + return rows + + +def write_rows(rows: list[dict[str, Any]], output: Path | None) -> None: + if not rows: + raise SystemExit("No rows to write.") + + fieldnames = sorted({key for row in rows for key in row}) + if output is None: + writer = csv.DictWriter( + sys.stdout, + fieldnames=fieldnames, + ) + writer.writeheader() + writer.writerows(rows) + return + + output.parent.mkdir(parents=True, exist_ok=True) + with output.open("w", encoding="utf-8", newline="") as file: + writer = csv.DictWriter(file, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + + +def main() -> int: + args = parse_args() + years = parse_years(args.years) + left = summarize_directory(args.left, years, profile=args.profile_support) + right = ( + summarize_directory(args.right, years, profile=args.profile_support) + if args.right is not None + else None + ) + rows = build_rows(left, right) + write_rows(rows, args.output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 047545b07b29149db60bd11c8628b499a2bb2809 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 31 Mar 2026 08:30:39 -0400 Subject: [PATCH 09/31] Add support concentration gates to calibration --- .../cps/long_term/calibration_profiles.py | 137 ++++++++++++++++++ .../test_long_term_calibration_contract.py | 55 +++++++ 2 files changed, 192 insertions(+) diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py index 49d13a612..7a5d45ea2 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py @@ -12,6 +12,10 @@ class ApproximateCalibrationWindow: max_age_error_pct: float max_negative_weight_pct: float | None = 0.0 age_bucket_size: int | None = None + min_positive_household_count: int | None = None + min_effective_sample_size: float | None = None + max_top_10_weight_share_pct: float | None = None + max_top_100_weight_share_pct: float | None = None def applies(self, year: int) -> bool: if year < self.start_year: @@ -35,6 +39,10 @@ class CalibrationProfile: max_constraint_error_pct: float = 0.1 max_age_error_pct: float = 0.1 max_negative_weight_pct: float | None = None + min_positive_household_count: int | None = None + min_effective_sample_size: float | None = None + max_top_10_weight_share_pct: float | None = None + max_top_100_weight_share_pct: float | None = None approximate_windows: tuple[ApproximateCalibrationWindow, ...] = field( default_factory=tuple ) @@ -51,6 +59,10 @@ def to_dict(self) -> dict[str, Any]: max_age_error_pct=0.5, max_negative_weight_pct=0.0, age_bucket_size=5, + min_positive_household_count=1000, + min_effective_sample_size=75.0, + max_top_10_weight_share_pct=25.0, + max_top_100_weight_share_pct=95.0, ), ApproximateCalibrationWindow( start_year=2079, @@ -59,6 +71,10 @@ def to_dict(self) -> dict[str, Any]: max_age_error_pct=10.0, max_negative_weight_pct=0.0, age_bucket_size=5, + min_positive_household_count=1000, + min_effective_sample_size=75.0, + max_top_10_weight_share_pct=25.0, + max_top_100_weight_share_pct=95.0, ), ApproximateCalibrationWindow( start_year=2086, @@ -67,6 +83,10 @@ def to_dict(self) -> dict[str, Any]: max_age_error_pct=20.0, max_negative_weight_pct=0.0, age_bucket_size=5, + min_positive_household_count=1000, + min_effective_sample_size=75.0, + max_top_10_weight_share_pct=25.0, + max_top_100_weight_share_pct=95.0, ), ApproximateCalibrationWindow( start_year=2096, @@ -75,6 +95,10 @@ def to_dict(self) -> dict[str, Any]: max_age_error_pct=35.0, max_negative_weight_pct=0.0, age_bucket_size=5, + min_positive_household_count=1000, + min_effective_sample_size=75.0, + max_top_10_weight_share_pct=25.0, + max_top_100_weight_share_pct=95.0, ), ) @@ -90,6 +114,10 @@ def to_dict(self) -> dict[str, Any]: use_h6_reform=False, use_tob=False, allow_greg_fallback=False, + min_positive_household_count=1000, + min_effective_sample_size=75.0, + max_top_10_weight_share_pct=25.0, + max_top_100_weight_share_pct=95.0, ), "ss": CalibrationProfile( name="ss", @@ -101,6 +129,10 @@ def to_dict(self) -> dict[str, Any]: use_h6_reform=False, use_tob=False, max_negative_weight_pct=0.0, + min_positive_household_count=1000, + min_effective_sample_size=75.0, + max_top_10_weight_share_pct=25.0, + max_top_100_weight_share_pct=95.0, approximate_windows=DEFAULT_LONG_RUN_APPROXIMATE_WINDOWS, ), "ss-payroll": CalibrationProfile( @@ -113,6 +145,10 @@ def to_dict(self) -> dict[str, Any]: use_h6_reform=False, use_tob=False, max_negative_weight_pct=0.0, + min_positive_household_count=1000, + min_effective_sample_size=75.0, + max_top_10_weight_share_pct=25.0, + max_top_100_weight_share_pct=95.0, approximate_windows=DEFAULT_LONG_RUN_APPROXIMATE_WINDOWS, ), "ss-payroll-tob": CalibrationProfile( @@ -125,6 +161,10 @@ def to_dict(self) -> dict[str, Any]: use_h6_reform=False, use_tob=True, max_negative_weight_pct=0.0, + min_positive_household_count=1000, + min_effective_sample_size=75.0, + max_top_10_weight_share_pct=25.0, + max_top_100_weight_share_pct=95.0, approximate_windows=DEFAULT_LONG_RUN_APPROXIMATE_WINDOWS, ), "ss-payroll-tob-h6": CalibrationProfile( @@ -137,6 +177,10 @@ def to_dict(self) -> dict[str, Any]: use_h6_reform=True, use_tob=True, max_negative_weight_pct=0.0, + min_positive_household_count=1000, + min_effective_sample_size=75.0, + max_top_10_weight_share_pct=25.0, + max_top_100_weight_share_pct=95.0, approximate_windows=DEFAULT_LONG_RUN_APPROXIMATE_WINDOWS, ), } @@ -249,12 +293,33 @@ def validate_calibration_audit( ) if quality == "exact": + window = approximate_window_for_year(profile, year) return _collect_threshold_issues( audit, profile, max_constraint_error_pct=profile.max_constraint_error_pct, max_age_error_pct=profile.max_age_error_pct, max_negative_weight_pct=profile.max_negative_weight_pct, + min_positive_household_count=( + window.min_positive_household_count + if window is not None + else profile.min_positive_household_count + ), + min_effective_sample_size=( + window.min_effective_sample_size + if window is not None + else profile.min_effective_sample_size + ), + max_top_10_weight_share_pct=( + window.max_top_10_weight_share_pct + if window is not None + else profile.max_top_10_weight_share_pct + ), + max_top_100_weight_share_pct=( + window.max_top_100_weight_share_pct + if window is not None + else profile.max_top_100_weight_share_pct + ), ) if quality == "approximate": @@ -266,6 +331,10 @@ def validate_calibration_audit( max_constraint_error_pct=profile.max_constraint_error_pct, max_age_error_pct=profile.max_age_error_pct, max_negative_weight_pct=profile.max_negative_weight_pct, + min_positive_household_count=profile.min_positive_household_count, + min_effective_sample_size=profile.min_effective_sample_size, + max_top_10_weight_share_pct=profile.max_top_10_weight_share_pct, + max_top_100_weight_share_pct=profile.max_top_100_weight_share_pct, ) issues.append( "Approximate calibration is not permitted for this profile/year" @@ -277,6 +346,10 @@ def validate_calibration_audit( max_constraint_error_pct=window.max_constraint_error_pct, max_age_error_pct=window.max_age_error_pct, max_negative_weight_pct=window.max_negative_weight_pct, + min_positive_household_count=window.min_positive_household_count, + min_effective_sample_size=window.min_effective_sample_size, + max_top_10_weight_share_pct=window.max_top_10_weight_share_pct, + max_top_100_weight_share_pct=window.max_top_100_weight_share_pct, ) exact_issues = _collect_threshold_issues( @@ -285,6 +358,10 @@ def validate_calibration_audit( max_constraint_error_pct=profile.max_constraint_error_pct, max_age_error_pct=profile.max_age_error_pct, max_negative_weight_pct=profile.max_negative_weight_pct, + min_positive_household_count=profile.min_positive_household_count, + min_effective_sample_size=profile.min_effective_sample_size, + max_top_10_weight_share_pct=profile.max_top_10_weight_share_pct, + max_top_100_weight_share_pct=profile.max_top_100_weight_share_pct, ) window = approximate_window_for_year(profile, year) if window is None: @@ -297,6 +374,10 @@ def validate_calibration_audit( max_constraint_error_pct=window.max_constraint_error_pct, max_age_error_pct=window.max_age_error_pct, max_negative_weight_pct=window.max_negative_weight_pct, + min_positive_household_count=window.min_positive_household_count, + min_effective_sample_size=window.min_effective_sample_size, + max_top_10_weight_share_pct=window.max_top_10_weight_share_pct, + max_top_100_weight_share_pct=window.max_top_100_weight_share_pct, ) return approximate_issues + [ "Calibration quality aggregate exceeds approximate thresholds" @@ -310,6 +391,10 @@ def _collect_threshold_issues( max_constraint_error_pct: float | None, max_age_error_pct: float | None, max_negative_weight_pct: float | None, + min_positive_household_count: int | None, + min_effective_sample_size: float | None, + max_top_10_weight_share_pct: float | None, + max_top_100_weight_share_pct: float | None, ) -> list[str]: issues: list[str] = [] @@ -346,6 +431,50 @@ def _collect_threshold_issues( f"{max_negative_weight_pct:.3f}%" ) + positive_count = audit.get("positive_weight_count") + if ( + min_positive_household_count is not None + and positive_count is not None + and positive_count < min_positive_household_count + ): + issues.append( + f"Positive household count {positive_count} is below " + f"{min_positive_household_count}" + ) + + ess = audit.get("effective_sample_size") + if ( + min_effective_sample_size is not None + and ess is not None + and ess < min_effective_sample_size + ): + issues.append( + f"Effective sample size {ess:.3f} is below " + f"{min_effective_sample_size:.3f}" + ) + + top_10_share = audit.get("top_10_weight_share_pct") + if ( + max_top_10_weight_share_pct is not None + and top_10_share is not None + and top_10_share > max_top_10_weight_share_pct + ): + issues.append( + f"Top-10 weight share {top_10_share:.3f}% exceeds " + f"{max_top_10_weight_share_pct:.3f}%" + ) + + top_100_share = audit.get("top_100_weight_share_pct") + if ( + max_top_100_weight_share_pct is not None + and top_100_share is not None + and top_100_share > max_top_100_weight_share_pct + ): + issues.append( + f"Top-100 weight share {top_100_share:.3f}% exceeds " + f"{max_top_100_weight_share_pct:.3f}%" + ) + return issues @@ -361,6 +490,10 @@ def classify_calibration_quality( max_constraint_error_pct=profile.max_constraint_error_pct, max_age_error_pct=profile.max_age_error_pct, max_negative_weight_pct=profile.max_negative_weight_pct, + min_positive_household_count=profile.min_positive_household_count, + min_effective_sample_size=profile.min_effective_sample_size, + max_top_10_weight_share_pct=profile.max_top_10_weight_share_pct, + max_top_100_weight_share_pct=profile.max_top_100_weight_share_pct, ) if not exact_issues: return "exact" @@ -377,6 +510,10 @@ def classify_calibration_quality( max_constraint_error_pct=window.max_constraint_error_pct, max_age_error_pct=window.max_age_error_pct, max_negative_weight_pct=window.max_negative_weight_pct, + min_positive_household_count=window.min_positive_household_count, + min_effective_sample_size=window.min_effective_sample_size, + max_top_10_weight_share_pct=window.max_top_10_weight_share_pct, + max_top_100_weight_share_pct=window.max_top_100_weight_share_pct, ) if not approximate_issues: return "approximate" diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index 8b23561c7..4d1fb8f2b 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -50,6 +50,10 @@ def test_named_profile_lookup(): assert profile.use_h6_reform is False assert profile.max_negative_weight_pct == 0.0 assert profile.approximate_windows[0].age_bucket_size == 5 + assert profile.min_positive_household_count == 1000 + assert profile.min_effective_sample_size == 75.0 + assert profile.max_top_10_weight_share_pct == 25.0 + assert profile.max_top_100_weight_share_pct == 95.0 def test_age_bin_helpers_preserve_population_totals(): @@ -193,6 +197,57 @@ def test_entropy_profile_rejects_negative_weights(): assert any("Negative weight share" in issue for issue in issues) +def test_support_thresholds_reject_concentrated_weights(): + profile = get_profile("ss-payroll-tob") + issues = validate_calibration_audit( + { + "fell_back_to_ipf": False, + "age_max_pct_error": 0.0, + "negative_weight_pct": 0.0, + "positive_weight_count": 90, + "effective_sample_size": 57.6, + "top_10_weight_share_pct": 26.6, + "top_100_weight_share_pct": 100.0, + "constraints": { + "ss_total": {"pct_error": 0.0}, + "payroll_total": {"pct_error": 0.0}, + "oasdi_tob": {"pct_error": 0.0}, + "hi_tob": {"pct_error": 0.0}, + }, + }, + profile, + year=2075, + quality="exact", + ) + assert any("Positive household count" in issue for issue in issues) + assert any("Top-10 weight share" in issue for issue in issues) + assert any("Top-100 weight share" in issue for issue in issues) + + +def test_classify_calibration_quality_marks_support_collapse_aggregate(): + profile = get_profile("ss-payroll-tob") + quality = classify_calibration_quality( + { + "fell_back_to_ipf": False, + "age_max_pct_error": 0.0, + "negative_weight_pct": 0.0, + "positive_weight_count": 6840, + "effective_sample_size": 24.98, + "top_10_weight_share_pct": 54.8, + "top_100_weight_share_pct": 97.4, + "constraints": { + "ss_total": {"pct_error": 0.0}, + "payroll_total": {"pct_error": 0.0}, + "oasdi_tob": {"pct_error": 0.0}, + "hi_tob": {"pct_error": 0.0}, + }, + }, + profile, + year=2075, + ) + assert quality == "aggregate" + + def test_approximate_window_is_year_bounded(): profile = get_profile("ss-payroll-tob") quality = classify_calibration_quality( From 5525b3d247a69aeff8e6233300439ef7893549e8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 31 Mar 2026 10:53:00 -0400 Subject: [PATCH 10/31] Add long-run support augmentation diagnostics --- .../evaluate_support_augmentation.py | 215 +++++++++ .../cps/long_term/support_augmentation.py | 427 ++++++++++++++++++ .../test_long_term_calibration_contract.py | 87 ++++ 3 files changed, 729 insertions(+) create mode 100644 policyengine_us_data/datasets/cps/long_term/evaluate_support_augmentation.py create mode 100644 policyengine_us_data/datasets/cps/long_term/support_augmentation.py diff --git a/policyengine_us_data/datasets/cps/long_term/evaluate_support_augmentation.py b/policyengine_us_data/datasets/cps/long_term/evaluate_support_augmentation.py new file mode 100644 index 000000000..de4ed2860 --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/evaluate_support_augmentation.py @@ -0,0 +1,215 @@ +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +from policyengine_us import Microsimulation + +from calibration import ( + _build_constraint_dataframe_and_controls, + assess_nonnegative_feasibility, +) +from calibration_profiles import approximate_window_for_year, get_profile +from projection_utils import ( + aggregate_age_targets, + aggregate_household_age_matrix, + build_age_bins, + build_household_age_matrix, +) +from ssa_data import ( + get_long_term_target_source, + load_hi_tob_projections, + load_oasdi_tob_projections, + load_ssa_age_projections, + load_ssa_benefit_projections, + load_taxable_payroll_projections, + set_long_term_target_source, +) +from support_augmentation import build_augmented_dataset + + +DEFAULT_DATASET = "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5" +BASE_YEAR = 2024 + + +def _evaluate_dataset( + *, + dataset: str | object, + dataset_label: str, + year: int, + profile_name: str, +) -> dict[str, object]: + profile = get_profile(profile_name) + sim = Microsimulation(dataset=dataset) + target_matrix = load_ssa_age_projections(start_year=year, end_year=year) + n_ages = target_matrix.shape[0] + X, _, _ = build_household_age_matrix(sim, n_ages=n_ages) + + approximate_window = approximate_window_for_year(profile, year) + age_bucket_size = ( + approximate_window.age_bucket_size + if approximate_window is not None + else None + ) + if age_bucket_size and age_bucket_size > 1: + age_bins = build_age_bins(n_ages=n_ages, bucket_size=age_bucket_size) + X_current = aggregate_household_age_matrix(X, age_bins) + y_target = aggregate_age_targets(target_matrix, age_bins)[:, 0] + else: + X_current = X + y_target = target_matrix[:, 0] + age_bucket_size = 1 + + household_series = sim.calculate("household_id", period=year, map_to="household") + baseline_weights = household_series.weights.values + + ss_values = None + ss_target = None + if profile.use_ss: + ss_values = sim.calculate("social_security", period=year, map_to="household").values + ss_target = load_ssa_benefit_projections(year) + + payroll_values = None + payroll_target = None + if profile.use_payroll: + payroll_values = ( + sim.calculate( + "taxable_earnings_for_social_security", + period=year, + map_to="household", + ).values + + sim.calculate( + "social_security_taxable_self_employment_income", + period=year, + map_to="household", + ).values + ) + payroll_target = load_taxable_payroll_projections(year) + + oasdi_tob_values = None + oasdi_tob_target = None + hi_tob_values = None + hi_tob_target = None + if profile.use_tob: + oasdi_tob_values = sim.calculate( + "tob_revenue_oasdi", + period=year, + map_to="household", + ).values + hi_tob_values = sim.calculate( + "tob_revenue_medicare_hi", + period=year, + map_to="household", + ).values + oasdi_tob_target = load_oasdi_tob_projections(year) + hi_tob_target = load_hi_tob_projections(year) + + aux_df, controls = _build_constraint_dataframe_and_controls( + X_current, + y_target, + ss_values=ss_values, + ss_target=ss_target, + payroll_values=payroll_values, + payroll_target=payroll_target, + oasdi_tob_values=oasdi_tob_values, + oasdi_tob_target=oasdi_tob_target, + hi_tob_values=hi_tob_values, + hi_tob_target=hi_tob_target, + n_ages=X_current.shape[1], + ) + targets = np.array(list(controls.values()), dtype=float) + feasibility = assess_nonnegative_feasibility( + aux_df.to_numpy(dtype=float), + targets, + ) + + return { + "dataset": dataset_label, + "year": year, + "profile": profile.name, + "target_source": get_long_term_target_source(), + "household_count": int(len(baseline_weights)), + "age_bucket_size": int(age_bucket_size), + "constraint_count": int(len(targets)), + "best_case_max_pct_error": feasibility["best_case_max_pct_error"], + "feasibility_status": feasibility["status"], + "feasibility_message": feasibility["message"], + } + + +def main() -> None: + parser = argparse.ArgumentParser( + description=( + "Compare late-year nonnegative feasibility before and after support augmentation." + ) + ) + parser.add_argument("year", type=int, help="Projection year to evaluate.") + parser.add_argument( + "--profile", + default="ss-payroll-tob", + help="Calibration profile to evaluate.", + ) + parser.add_argument( + "--target-source", + default="trustees_2025_current_law", + help="Named long-run target source package.", + ) + parser.add_argument( + "--dataset", + default=DEFAULT_DATASET, + help="Base dataset path or HF reference.", + ) + parser.add_argument( + "--support-augmentation", + default="late-clone-v1", + help="Support augmentation profile name.", + ) + args = parser.parse_args() + + set_long_term_target_source(args.target_source) + + base_result = _evaluate_dataset( + dataset=args.dataset, + dataset_label="base", + year=args.year, + profile_name=args.profile, + ) + augmented_dataset, augmentation_report = build_augmented_dataset( + base_dataset=args.dataset, + base_year=BASE_YEAR, + profile=args.support_augmentation, + ) + augmented_result = _evaluate_dataset( + dataset=augmented_dataset, + dataset_label=args.support_augmentation, + year=args.year, + profile_name=args.profile, + ) + + report = { + "year": args.year, + "profile": args.profile, + "target_source": args.target_source, + "augmentation": augmentation_report, + "results": { + "base": base_result, + "augmented": augmented_result, + "delta_best_case_max_pct_error": ( + None + if base_result["best_case_max_pct_error"] is None + or augmented_result["best_case_max_pct_error"] is None + else augmented_result["best_case_max_pct_error"] + - base_result["best_case_max_pct_error"] + ), + "delta_household_count": ( + augmented_result["household_count"] - base_result["household_count"] + ), + }, + } + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/policyengine_us_data/datasets/cps/long_term/support_augmentation.py b/policyengine_us_data/datasets/cps/long_term/support_augmentation.py new file mode 100644 index 000000000..3503ab0be --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/support_augmentation.py @@ -0,0 +1,427 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Literal + +import numpy as np +import pandas as pd +from policyengine_core.data.dataset import Dataset +from policyengine_us import Microsimulation + + +SS_COMPONENTS = ( + "social_security_retirement", + "social_security_disability", + "social_security_survivors", + "social_security_dependents", +) +PAYROLL_COMPONENTS = ( + "employment_income_before_lsr", + "self_employment_income_before_lsr", +) +ENTITY_ID_COLUMNS = { + "household": ("household_id", "person_household_id"), + "family": ("family_id", "person_family_id"), + "tax_unit": ("tax_unit_id", "person_tax_unit_id"), + "spm_unit": ("spm_unit_id", "person_spm_unit_id"), + "marital_unit": ("marital_unit_id", "person_marital_unit_id"), +} +PERSON_ID_COLUMN = "person_id" + + +ConstraintState = Literal["any", "positive", "nonpositive"] + + +@dataclass(frozen=True) +class AgeShiftCloneRule: + name: str + min_max_age: int + max_max_age: int + age_shift: int + ss_state: ConstraintState = "any" + payroll_state: ConstraintState = "any" + clone_weight_scale: float = 0.25 + + +@dataclass(frozen=True) +class SupportAugmentationProfile: + name: str + description: str + rules: tuple[AgeShiftCloneRule, ...] + + +LATE_CLONE_V1 = SupportAugmentationProfile( + name="late-clone-v1", + description=( + "Age-shifted donor households to expand late-year support for older " + "beneficiary, older beneficiary-plus-payroll, and payroll-only households." + ), + rules=( + AgeShiftCloneRule( + name="ss_only_65_74_to_75_84", + min_max_age=65, + max_max_age=74, + age_shift=10, + ss_state="positive", + payroll_state="nonpositive", + clone_weight_scale=0.35, + ), + AgeShiftCloneRule( + name="ss_only_75_84_to_85_plus", + min_max_age=75, + max_max_age=84, + age_shift=10, + ss_state="positive", + payroll_state="nonpositive", + clone_weight_scale=0.5, + ), + AgeShiftCloneRule( + name="ss_pay_65_74_to_75_84", + min_max_age=65, + max_max_age=74, + age_shift=10, + ss_state="positive", + payroll_state="positive", + clone_weight_scale=0.35, + ), + AgeShiftCloneRule( + name="ss_pay_75_84_to_85_plus", + min_max_age=75, + max_max_age=84, + age_shift=10, + ss_state="positive", + payroll_state="positive", + clone_weight_scale=0.5, + ), + AgeShiftCloneRule( + name="pay_only_55_64_to_65_74", + min_max_age=55, + max_max_age=64, + age_shift=10, + ss_state="nonpositive", + payroll_state="positive", + clone_weight_scale=0.2, + ), + ), +) + +LATE_CLONE_V2 = SupportAugmentationProfile( + name="late-clone-v2", + description=( + "More aggressive age-shifted donor households that test whether the " + "late-year infeasibility is driven by missing older payroll-rich support." + ), + rules=( + *LATE_CLONE_V1.rules, + AgeShiftCloneRule( + name="pay_only_45_64_to_75_84", + min_max_age=45, + max_max_age=64, + age_shift=20, + ss_state="nonpositive", + payroll_state="positive", + clone_weight_scale=0.15, + ), + AgeShiftCloneRule( + name="pay_only_55_64_to_85_plus", + min_max_age=55, + max_max_age=64, + age_shift=30, + ss_state="nonpositive", + payroll_state="positive", + clone_weight_scale=0.1, + ), + AgeShiftCloneRule( + name="ss_pay_65_74_to_85_plus", + min_max_age=65, + max_max_age=74, + age_shift=20, + ss_state="positive", + payroll_state="positive", + clone_weight_scale=0.2, + ), + ), +) + + +NAMED_SUPPORT_AUGMENTATION_PROFILES = { + LATE_CLONE_V1.name: LATE_CLONE_V1, + LATE_CLONE_V2.name: LATE_CLONE_V2, +} + + +def _period_column(name: str, base_year: int) -> str: + return f"{name}__{base_year}" + + +def get_support_augmentation_profile(name: str) -> SupportAugmentationProfile: + try: + return NAMED_SUPPORT_AUGMENTATION_PROFILES[name] + except KeyError as error: + valid = ", ".join(sorted(NAMED_SUPPORT_AUGMENTATION_PROFILES)) + raise ValueError( + f"Unknown support augmentation profile '{name}'. Valid profiles: {valid}" + ) from error + + +def household_support_summary( + input_df: pd.DataFrame, + *, + base_year: int, +) -> pd.DataFrame: + household_id_col = _period_column("household_id", base_year) + household_weight_col = _period_column("household_weight", base_year) + age_col = _period_column("age", base_year) + + required = [household_id_col, household_weight_col, age_col] + required.extend(_period_column(component, base_year) for component in SS_COMPONENTS) + required.extend( + _period_column(component, base_year) for component in PAYROLL_COMPONENTS + ) + missing = [column for column in required if column not in input_df.columns] + if missing: + raise ValueError( + "Input dataframe is missing required support columns: " + + ", ".join(sorted(missing)) + ) + + aggregations: dict[str, str] = { + age_col: "max", + household_weight_col: "max", + } + aggregations.update( + {_period_column(component, base_year): "sum" for component in SS_COMPONENTS} + ) + aggregations.update( + {_period_column(component, base_year): "sum" for component in PAYROLL_COMPONENTS} + ) + + summary = ( + input_df.groupby(household_id_col, sort=False) + .agg(aggregations) + .rename( + columns={ + age_col: "max_age", + household_weight_col: "baseline_weight", + } + ) + ) + summary["ss_total"] = summary[ + [_period_column(component, base_year) for component in SS_COMPONENTS] + ].sum(axis=1) + summary["payroll_total"] = summary[ + [_period_column(component, base_year) for component in PAYROLL_COMPONENTS] + ].sum(axis=1) + return summary + + +def _match_state(values: pd.Series, state: ConstraintState) -> pd.Series: + if state == "any": + return pd.Series(True, index=values.index) + if state == "positive": + return values > 0 + if state == "nonpositive": + return values <= 0 + raise ValueError(f"Unsupported state '{state}'") + + +def select_donor_households( + summary: pd.DataFrame, + rule: AgeShiftCloneRule, +) -> pd.Index: + age_mask = summary["max_age"].between(rule.min_max_age, rule.max_max_age) + positive_weight_mask = summary["baseline_weight"] > 0 + ss_mask = _match_state(summary["ss_total"], rule.ss_state) + payroll_mask = _match_state(summary["payroll_total"], rule.payroll_state) + return summary.index[age_mask & positive_weight_mask & ss_mask & payroll_mask] + + +def _next_entity_id(values: pd.Series) -> int: + non_null = values.dropna() + if non_null.empty: + return 1 + return int(non_null.max()) + 1 + + +def _cast_mapped_ids(series: pd.Series, mapped: pd.Series) -> pd.Series: + dtype = series.dtype + if pd.api.types.is_integer_dtype(dtype): + return mapped.astype(dtype) + if pd.api.types.is_float_dtype(dtype): + return mapped.astype(dtype) + return mapped + + +def clone_households_with_age_shift( + input_df: pd.DataFrame, + *, + base_year: int, + household_ids: pd.Index, + age_shift: int, + clone_weight_scale: float, + id_counters: dict[str, int] | None = None, +) -> tuple[pd.DataFrame, dict[str, int]]: + if household_ids.empty: + return input_df.iloc[0:0].copy(), ( + id_counters.copy() if id_counters is not None else {} + ) + + household_id_col = _period_column("household_id", base_year) + person_id_col = _period_column(PERSON_ID_COLUMN, base_year) + age_col = _period_column("age", base_year) + household_weight_col = _period_column("household_weight", base_year) + person_weight_col = _period_column("person_weight", base_year) + + donors = input_df[input_df[household_id_col].isin(household_ids)].copy() + + next_ids = ( + id_counters.copy() + if id_counters is not None + else { + entity_name: _next_entity_id( + input_df[_period_column(columns[0], base_year)] + ) + for entity_name, columns in ENTITY_ID_COLUMNS.items() + } + ) + if "person" not in next_ids: + next_ids["person"] = _next_entity_id(input_df[person_id_col]) + + household_map = { + original_id: next_ids["household"] + offset + for offset, original_id in enumerate( + pd.unique(donors[household_id_col].dropna()) + ) + } + next_ids["household"] += len(household_map) + + for entity_name, columns in ENTITY_ID_COLUMNS.items(): + column = _period_column(columns[0], base_year) + if entity_name == "household" or column not in donors.columns: + continue + unique_ids = pd.unique(donors[column].dropna()) + mapping = { + original_id: next_ids[entity_name] + offset + for offset, original_id in enumerate(unique_ids) + } + next_ids[entity_name] += len(mapping) + for raw_column in columns: + mapped_column = _period_column(raw_column, base_year) + if mapped_column not in donors.columns: + continue + mapped = donors[mapped_column].map(mapping) + donors[mapped_column] = _cast_mapped_ids(donors[mapped_column], mapped) + + person_map = { + original_id: next_ids["person"] + offset + for offset, original_id in enumerate(pd.unique(donors[person_id_col].dropna())) + } + next_ids["person"] += len(person_map) + donors[person_id_col] = _cast_mapped_ids( + donors[person_id_col], donors[person_id_col].map(person_map) + ) + + for raw_column in ENTITY_ID_COLUMNS["household"]: + mapped_column = _period_column(raw_column, base_year) + donors[mapped_column] = _cast_mapped_ids( + donors[mapped_column], donors[mapped_column].map(household_map) + ) + + donors[age_col] = np.minimum(donors[age_col].astype(float) + age_shift, 85) + + if household_weight_col in donors.columns: + donors[household_weight_col] = ( + donors[household_weight_col].astype(float) * clone_weight_scale + ) + if person_weight_col in donors.columns: + donors[person_weight_col] = ( + donors[person_weight_col].astype(float) * clone_weight_scale + ) + + return donors, next_ids + + +def augment_input_dataframe( + input_df: pd.DataFrame, + *, + base_year: int, + profile: str | SupportAugmentationProfile, +) -> tuple[pd.DataFrame, dict[str, Any]]: + profile_obj = ( + get_support_augmentation_profile(profile) + if isinstance(profile, str) + else profile + ) + summary = household_support_summary(input_df, base_year=base_year) + + clone_frames = [] + rule_reports: list[dict[str, Any]] = [] + id_counters = { + entity_name: _next_entity_id(input_df[_period_column(columns[0], base_year)]) + for entity_name, columns in ENTITY_ID_COLUMNS.items() + } + id_counters["person"] = _next_entity_id( + input_df[_period_column(PERSON_ID_COLUMN, base_year)] + ) + for rule in profile_obj.rules: + donor_households = select_donor_households(summary, rule) + clone_df, id_counters = clone_households_with_age_shift( + input_df, + base_year=base_year, + household_ids=donor_households, + age_shift=rule.age_shift, + clone_weight_scale=rule.clone_weight_scale, + id_counters=id_counters, + ) + clone_frames.append(clone_df) + rule_reports.append( + { + "rule": rule.name, + "donor_household_count": int(len(donor_households)), + "clone_household_count": int( + clone_df[_period_column("household_id", base_year)].nunique() + ) + if not clone_df.empty + else 0, + "clone_person_count": int(len(clone_df)), + "age_shift": rule.age_shift, + "clone_weight_scale": rule.clone_weight_scale, + } + ) + + if clone_frames: + augmented_df = pd.concat([input_df, *clone_frames], ignore_index=True) + else: + augmented_df = input_df.copy() + + report = { + "profile": profile_obj.name, + "description": profile_obj.description, + "base_household_count": int( + input_df[_period_column("household_id", base_year)].nunique() + ), + "base_person_count": int(len(input_df)), + "augmented_household_count": int( + augmented_df[_period_column("household_id", base_year)].nunique() + ), + "augmented_person_count": int(len(augmented_df)), + "rules": rule_reports, + } + return augmented_df, report + + +def build_augmented_dataset( + *, + base_dataset: str, + base_year: int, + profile: str | SupportAugmentationProfile, +) -> tuple[Dataset, dict[str, Any]]: + sim = Microsimulation(dataset=base_dataset) + input_df = sim.to_input_dataframe() + augmented_df, report = augment_input_dataframe( + input_df, + base_year=base_year, + profile=profile, + ) + report["base_dataset"] = base_dataset + return Dataset.from_dataframe(augmented_df, base_year), report diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index 4d1fb8f2b..936e238bc 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -33,6 +33,13 @@ describe_long_term_target_source, load_taxable_payroll_projections, ) +from policyengine_us_data.datasets.cps.long_term.support_augmentation import ( + AgeShiftCloneRule, + SupportAugmentationProfile, + augment_input_dataframe, + household_support_summary, + select_donor_households, +) class ExplodingCalibrator: @@ -40,6 +47,35 @@ def calibrate(self, **kwargs): raise RuntimeError("boom") +def _toy_support_dataframe(): + return json.loads( + json.dumps( + { + "person_id__2024": [101, 102, 201, 202, 301], + "household_id__2024": [1, 1, 2, 2, 3], + "person_household_id__2024": [1, 1, 2, 2, 3], + "family_id__2024": [11.0, 11.0, 21.0, 21.0, 31.0], + "person_family_id__2024": [11, 11, 21, 21, 31], + "tax_unit_id__2024": [101, 101, 201, 202, 301], + "person_tax_unit_id__2024": [101, 101, 201, 202, 301], + "spm_unit_id__2024": [1001, 1001, 2001, 2001, 3001], + "person_spm_unit_id__2024": [1001, 1001, 2001, 2001, 3001], + "marital_unit_id__2024": [501, 501, 601, 602, 701], + "person_marital_unit_id__2024": [501, 501, 601, 602, 701], + "age__2024": [70.0, 68.0, 80.0, 77.0, 60.0], + "household_weight__2024": [10.0, 10.0, 8.0, 8.0, 5.0], + "person_weight__2024": [10.0, 10.0, 8.0, 8.0, 5.0], + "social_security_retirement__2024": [20_000.0, 0.0, 30_000.0, 0.0, 0.0], + "social_security_disability__2024": [0.0, 0.0, 0.0, 0.0, 0.0], + "social_security_survivors__2024": [0.0, 0.0, 0.0, 0.0, 0.0], + "social_security_dependents__2024": [0.0, 0.0, 0.0, 0.0, 0.0], + "employment_income_before_lsr__2024": [5_000.0, 0.0, 12_000.0, 0.0, 50_000.0], + "self_employment_income_before_lsr__2024": [0.0, 0.0, 0.0, 0.0, 0.0], + } + ) + ) + + def test_named_profile_lookup(): profile = get_profile("ss-payroll-tob") assert profile.calibration_method == "entropy" @@ -56,6 +92,57 @@ def test_named_profile_lookup(): assert profile.max_top_100_weight_share_pct == 95.0 +def test_support_augmentation_selects_expected_donors(): + import pandas as pd + + df = pd.DataFrame(_toy_support_dataframe()) + summary = household_support_summary(df, base_year=2024) + rule = AgeShiftCloneRule( + name="older_ss_pay", + min_max_age=65, + max_max_age=74, + age_shift=10, + ss_state="positive", + payroll_state="positive", + ) + donors = select_donor_households(summary, rule) + assert list(donors) == [1] + + +def test_support_augmentation_clones_households_with_new_ids(): + import pandas as pd + + df = pd.DataFrame(_toy_support_dataframe()) + profile = SupportAugmentationProfile( + name="test-profile", + description="Toy support augmentation profile.", + rules=( + AgeShiftCloneRule( + name="older_ss_pay", + min_max_age=65, + max_max_age=74, + age_shift=10, + ss_state="positive", + payroll_state="positive", + clone_weight_scale=0.5, + ), + ), + ) + augmented_df, report = augment_input_dataframe( + df, + base_year=2024, + profile=profile, + ) + assert report["base_household_count"] == 3 + assert report["augmented_household_count"] == 4 + cloned_household_ids = set(augmented_df["household_id__2024"].unique()) - {1, 2, 3} + assert len(cloned_household_ids) == 1 + cloned_rows = augmented_df[augmented_df["household_id__2024"].isin(cloned_household_ids)] + assert cloned_rows["age__2024"].max() == pytest.approx(80.0) + assert cloned_rows["household_weight__2024"].iloc[0] == pytest.approx(5.0) + assert cloned_rows["person_id__2024"].min() > df["person_id__2024"].max() + + def test_age_bin_helpers_preserve_population_totals(): bins = build_age_bins(n_ages=86, bucket_size=5) assert bins[0] == (0, 5) From e3d9912147bec9e97c533690dabcd9680e94b869 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 31 Mar 2026 10:59:30 -0400 Subject: [PATCH 11/31] Probe composite long-run support augmentation --- .../cps/long_term/support_augmentation.py | 356 ++++++++++++++++-- .../test_long_term_calibration_contract.py | 39 ++ 2 files changed, 374 insertions(+), 21 deletions(-) diff --git a/policyengine_us_data/datasets/cps/long_term/support_augmentation.py b/policyengine_us_data/datasets/cps/long_term/support_augmentation.py index 3503ab0be..ae64c6dc4 100644 --- a/policyengine_us_data/datasets/cps/long_term/support_augmentation.py +++ b/policyengine_us_data/datasets/cps/long_term/support_augmentation.py @@ -19,6 +19,9 @@ "employment_income_before_lsr", "self_employment_income_before_lsr", ) +PAYROLL_TRANSFER_COMPONENTS = PAYROLL_COMPONENTS + ( + "w2_wages_from_qualified_business", +) ENTITY_ID_COLUMNS = { "household": ("household_id", "person_household_id"), "family": ("family_id", "person_family_id"), @@ -43,11 +46,26 @@ class AgeShiftCloneRule: clone_weight_scale: float = 0.25 +@dataclass(frozen=True) +class CompositePayrollRule: + name: str + recipient_min_max_age: int + recipient_max_max_age: int + donor_min_max_age: int + donor_max_max_age: int + recipient_ss_state: ConstraintState = "positive" + recipient_payroll_state: ConstraintState = "nonpositive" + donor_ss_state: ConstraintState = "nonpositive" + donor_payroll_state: ConstraintState = "positive" + payroll_transfer_scale: float = 1.0 + clone_weight_scale: float = 0.25 + + @dataclass(frozen=True) class SupportAugmentationProfile: name: str description: str - rules: tuple[AgeShiftCloneRule, ...] + rules: tuple[AgeShiftCloneRule | CompositePayrollRule, ...] LATE_CLONE_V1 = SupportAugmentationProfile( @@ -143,10 +161,110 @@ class SupportAugmentationProfile: ), ) +LATE_COMPOSITE_V1 = SupportAugmentationProfile( + name="late-composite-v1", + description=( + "Composite synthetic households that preserve older beneficiary age/SS " + "structure while injecting payroll from younger payroll-rich donors." + ), + rules=( + CompositePayrollRule( + name="ss_only_75_84_plus_payroll_from_55_64", + recipient_min_max_age=75, + recipient_max_max_age=84, + donor_min_max_age=55, + donor_max_max_age=64, + recipient_ss_state="positive", + recipient_payroll_state="nonpositive", + donor_ss_state="nonpositive", + donor_payroll_state="positive", + payroll_transfer_scale=1.0, + clone_weight_scale=0.2, + ), + CompositePayrollRule( + name="ss_only_85_plus_plus_payroll_from_55_64", + recipient_min_max_age=85, + recipient_max_max_age=85, + donor_min_max_age=55, + donor_max_max_age=64, + recipient_ss_state="positive", + recipient_payroll_state="nonpositive", + donor_ss_state="nonpositive", + donor_payroll_state="positive", + payroll_transfer_scale=0.75, + clone_weight_scale=0.15, + ), + CompositePayrollRule( + name="ss_pay_75_84_boost_from_45_64", + recipient_min_max_age=75, + recipient_max_max_age=84, + donor_min_max_age=45, + donor_max_max_age=64, + recipient_ss_state="positive", + recipient_payroll_state="positive", + donor_ss_state="nonpositive", + donor_payroll_state="positive", + payroll_transfer_scale=0.5, + clone_weight_scale=0.15, + ), + ), +) + +LATE_COMPOSITE_V2 = SupportAugmentationProfile( + name="late-composite-v2", + description=( + "Extreme composite synthetic households for diagnosing whether the late " + "frontier is limited by missing older payroll intensity." + ), + rules=( + CompositePayrollRule( + name="ss_only_75_84_plus_heavy_payroll_from_45_64", + recipient_min_max_age=75, + recipient_max_max_age=84, + donor_min_max_age=45, + donor_max_max_age=64, + recipient_ss_state="positive", + recipient_payroll_state="nonpositive", + donor_ss_state="nonpositive", + donor_payroll_state="positive", + payroll_transfer_scale=3.0, + clone_weight_scale=0.15, + ), + CompositePayrollRule( + name="ss_only_85_plus_heavy_payroll_from_45_64", + recipient_min_max_age=85, + recipient_max_max_age=85, + donor_min_max_age=45, + donor_max_max_age=64, + recipient_ss_state="positive", + recipient_payroll_state="nonpositive", + donor_ss_state="nonpositive", + donor_payroll_state="positive", + payroll_transfer_scale=2.0, + clone_weight_scale=0.1, + ), + CompositePayrollRule( + name="ss_pay_75_84_heavy_boost_from_45_64", + recipient_min_max_age=75, + recipient_max_max_age=84, + donor_min_max_age=45, + donor_max_max_age=64, + recipient_ss_state="positive", + recipient_payroll_state="positive", + donor_ss_state="nonpositive", + donor_payroll_state="positive", + payroll_transfer_scale=1.5, + clone_weight_scale=0.1, + ), + ), +) + NAMED_SUPPORT_AUGMENTATION_PROFILES = { LATE_CLONE_V1.name: LATE_CLONE_V1, LATE_CLONE_V2.name: LATE_CLONE_V2, + LATE_COMPOSITE_V1.name: LATE_COMPOSITE_V1, + LATE_COMPOSITE_V2.name: LATE_COMPOSITE_V2, } @@ -236,6 +354,21 @@ def select_donor_households( return summary.index[age_mask & positive_weight_mask & ss_mask & payroll_mask] +def select_households_for_composite_rule( + summary: pd.DataFrame, + *, + min_max_age: int, + max_max_age: int, + ss_state: ConstraintState, + payroll_state: ConstraintState, +) -> pd.Index: + age_mask = summary["max_age"].between(min_max_age, max_max_age) + positive_weight_mask = summary["baseline_weight"] > 0 + ss_mask = _match_state(summary["ss_total"], ss_state) + payroll_mask = _match_state(summary["payroll_total"], payroll_state) + return summary.index[age_mask & positive_weight_mask & ss_mask & payroll_mask] + + def _next_entity_id(values: pd.Series) -> int: non_null = values.dropna() if non_null.empty: @@ -341,6 +474,175 @@ def clone_households_with_age_shift( return donors, next_ids +def _household_component_totals( + input_df: pd.DataFrame, + *, + base_year: int, + components: tuple[str, ...], +) -> pd.DataFrame: + household_id_col = _period_column("household_id", base_year) + available = [ + _period_column(component, base_year) + for component in components + if _period_column(component, base_year) in input_df.columns + ] + if not available: + return pd.DataFrame(index=pd.Index([], dtype=int)) + return input_df.groupby(household_id_col, sort=False)[available].sum() + + +def _quantile_pair_households( + recipient_ids: pd.Index, + donor_ids: pd.Index, + summary: pd.DataFrame, +) -> list[tuple[int, int]]: + if recipient_ids.empty or donor_ids.empty: + return [] + recipient_order = summary.loc[recipient_ids].sort_values( + ["ss_total", "baseline_weight", "max_age"] + ).index.to_list() + donor_order = summary.loc[donor_ids].sort_values( + ["payroll_total", "baseline_weight", "max_age"] + ).index.to_list() + if len(donor_order) == 1: + donor_positions = np.zeros(len(recipient_order), dtype=int) + else: + donor_positions = np.linspace( + 0, + len(donor_order) - 1, + num=len(recipient_order), + ).round().astype(int) + return [ + (int(recipient_household_id), int(donor_order[position])) + for recipient_household_id, position in zip(recipient_order, donor_positions) + ] + + +def _select_payroll_target_row( + household_rows: pd.DataFrame, + *, + base_year: int, +) -> Any: + age_col = _period_column("age", base_year) + employment_col = _period_column("employment_income_before_lsr", base_year) + self_employment_col = _period_column("self_employment_income_before_lsr", base_year) + adults = household_rows[household_rows[age_col] >= 18] + if adults.empty: + adults = household_rows + existing_payroll = ( + adults.get(employment_col, 0).astype(float) + + adults.get(self_employment_col, 0).astype(float) + ) + if existing_payroll.gt(0).any(): + return existing_payroll.idxmax() + return adults[age_col].astype(float).idxmax() + + +def synthesize_composite_households( + input_df: pd.DataFrame, + *, + base_year: int, + summary: pd.DataFrame, + rule: CompositePayrollRule, + id_counters: dict[str, int] | None = None, +) -> tuple[pd.DataFrame, dict[str, int], dict[str, Any]]: + recipient_ids = select_households_for_composite_rule( + summary, + min_max_age=rule.recipient_min_max_age, + max_max_age=rule.recipient_max_max_age, + ss_state=rule.recipient_ss_state, + payroll_state=rule.recipient_payroll_state, + ) + donor_ids = select_households_for_composite_rule( + summary, + min_max_age=rule.donor_min_max_age, + max_max_age=rule.donor_max_max_age, + ss_state=rule.donor_ss_state, + payroll_state=rule.donor_payroll_state, + ) + recipient_pairs = _quantile_pair_households(recipient_ids, donor_ids, summary) + if not recipient_pairs: + return ( + input_df.iloc[0:0].copy(), + id_counters.copy() if id_counters is not None else {}, + { + "rule": rule.name, + "recipient_household_count": 0, + "donor_household_count": int(len(donor_ids)), + "composite_household_count": 0, + "composite_person_count": 0, + "payroll_transfer_scale": rule.payroll_transfer_scale, + }, + ) + + clone_df, next_ids = clone_households_with_age_shift( + input_df, + base_year=base_year, + household_ids=pd.Index([recipient for recipient, _ in recipient_pairs]), + age_shift=0, + clone_weight_scale=rule.clone_weight_scale, + id_counters=id_counters, + ) + household_id_col = _period_column("household_id", base_year) + original_household_col = _period_column("person_household_id", base_year) + payroll_totals = _household_component_totals( + input_df, + base_year=base_year, + components=PAYROLL_TRANSFER_COMPONENTS, + ) + + original_recipients = pd.unique( + input_df[input_df[household_id_col].isin([recipient for recipient, _ in recipient_pairs])][ + household_id_col + ] + ) + cloned_household_ids = pd.unique(clone_df[household_id_col]) + cloned_mapping = { + int(original): int(cloned) + for original, cloned in zip(original_recipients, cloned_household_ids) + } + + employment_col = _period_column("employment_income_before_lsr", base_year) + self_employment_col = _period_column("self_employment_income_before_lsr", base_year) + qbi_col = _period_column("w2_wages_from_qualified_business", base_year) + + for recipient_household_id, donor_household_id in recipient_pairs: + cloned_household_id = cloned_mapping[int(recipient_household_id)] + mask = clone_df[household_id_col] == cloned_household_id + target_row = _select_payroll_target_row( + clone_df.loc[mask], + base_year=base_year, + ) + donor_row = payroll_totals.loc[int(donor_household_id)] + clone_df.loc[target_row, employment_col] = ( + clone_df.loc[target_row, employment_col] + + float(donor_row.get(employment_col, 0.0)) * rule.payroll_transfer_scale + ) + clone_df.loc[target_row, self_employment_col] = ( + clone_df.loc[target_row, self_employment_col] + + float(donor_row.get(self_employment_col, 0.0)) + * rule.payroll_transfer_scale + ) + if qbi_col in clone_df.columns and qbi_col in donor_row.index: + clone_df.loc[target_row, qbi_col] = ( + clone_df.loc[target_row, qbi_col] + + float(donor_row.get(qbi_col, 0.0)) * rule.payroll_transfer_scale + ) + + return ( + clone_df, + next_ids, + { + "rule": rule.name, + "recipient_household_count": int(len(recipient_ids)), + "donor_household_count": int(len(donor_ids)), + "composite_household_count": int(clone_df[household_id_col].nunique()), + "composite_person_count": int(len(clone_df)), + "payroll_transfer_scale": rule.payroll_transfer_scale, + }, + ) + + def augment_input_dataframe( input_df: pd.DataFrame, *, @@ -364,30 +666,42 @@ def augment_input_dataframe( input_df[_period_column(PERSON_ID_COLUMN, base_year)] ) for rule in profile_obj.rules: - donor_households = select_donor_households(summary, rule) - clone_df, id_counters = clone_households_with_age_shift( + if isinstance(rule, AgeShiftCloneRule): + donor_households = select_donor_households(summary, rule) + clone_df, id_counters = clone_households_with_age_shift( + input_df, + base_year=base_year, + household_ids=donor_households, + age_shift=rule.age_shift, + clone_weight_scale=rule.clone_weight_scale, + id_counters=id_counters, + ) + clone_frames.append(clone_df) + rule_reports.append( + { + "rule": rule.name, + "donor_household_count": int(len(donor_households)), + "clone_household_count": int( + clone_df[_period_column("household_id", base_year)].nunique() + ) + if not clone_df.empty + else 0, + "clone_person_count": int(len(clone_df)), + "age_shift": rule.age_shift, + "clone_weight_scale": rule.clone_weight_scale, + } + ) + continue + + composite_df, id_counters, composite_report = synthesize_composite_households( input_df, base_year=base_year, - household_ids=donor_households, - age_shift=rule.age_shift, - clone_weight_scale=rule.clone_weight_scale, + summary=summary, + rule=rule, id_counters=id_counters, ) - clone_frames.append(clone_df) - rule_reports.append( - { - "rule": rule.name, - "donor_household_count": int(len(donor_households)), - "clone_household_count": int( - clone_df[_period_column("household_id", base_year)].nunique() - ) - if not clone_df.empty - else 0, - "clone_person_count": int(len(clone_df)), - "age_shift": rule.age_shift, - "clone_weight_scale": rule.clone_weight_scale, - } - ) + clone_frames.append(composite_df) + rule_reports.append(composite_report) if clone_frames: augmented_df = pd.concat([input_df, *clone_frames], ignore_index=True) diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index 936e238bc..1f988c1fb 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -35,6 +35,7 @@ ) from policyengine_us_data.datasets.cps.long_term.support_augmentation import ( AgeShiftCloneRule, + CompositePayrollRule, SupportAugmentationProfile, augment_input_dataframe, household_support_summary, @@ -143,6 +144,44 @@ def test_support_augmentation_clones_households_with_new_ids(): assert cloned_rows["person_id__2024"].min() > df["person_id__2024"].max() +def test_support_augmentation_synthesizes_composite_payroll_household(): + import pandas as pd + + df = pd.DataFrame(_toy_support_dataframe()) + profile = SupportAugmentationProfile( + name="composite-profile", + description="Toy composite support augmentation profile.", + rules=( + CompositePayrollRule( + name="older_ss_only_plus_payroll", + recipient_min_max_age=75, + recipient_max_max_age=84, + donor_min_max_age=55, + donor_max_max_age=64, + recipient_ss_state="positive", + recipient_payroll_state="positive", + donor_ss_state="nonpositive", + donor_payroll_state="positive", + payroll_transfer_scale=0.5, + clone_weight_scale=0.25, + ), + ), + ) + augmented_df, report = augment_input_dataframe( + df, + base_year=2024, + profile=profile, + ) + assert report["base_household_count"] == 3 + assert report["augmented_household_count"] == 4 + cloned_household_ids = set(augmented_df["household_id__2024"].unique()) - {1, 2, 3} + assert len(cloned_household_ids) == 1 + cloned_rows = augmented_df[augmented_df["household_id__2024"].isin(cloned_household_ids)] + assert cloned_rows["age__2024"].max() == pytest.approx(80.0) + assert cloned_rows["social_security_retirement__2024"].sum() == pytest.approx(30_000.0) + assert cloned_rows["employment_income_before_lsr__2024"].sum() == pytest.approx(37_000.0) + + def test_age_bin_helpers_preserve_population_totals(): bins = build_age_bins(n_ages=86, bucket_size=5) assert bins[0] == (0, 5) From 5b91f1e5637b40cab5b14773318a53ea85d9124a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 31 Mar 2026 12:47:29 -0400 Subject: [PATCH 12/31] Test appended synthetic late-year support --- .../cps/long_term/support_augmentation.py | 279 +++++++++++++++++- .../test_long_term_calibration_contract.py | 63 ++++ 2 files changed, 336 insertions(+), 6 deletions(-) diff --git a/policyengine_us_data/datasets/cps/long_term/support_augmentation.py b/policyengine_us_data/datasets/cps/long_term/support_augmentation.py index ae64c6dc4..5ce817cae 100644 --- a/policyengine_us_data/datasets/cps/long_term/support_augmentation.py +++ b/policyengine_us_data/datasets/cps/long_term/support_augmentation.py @@ -61,11 +61,31 @@ class CompositePayrollRule: clone_weight_scale: float = 0.25 +@dataclass(frozen=True) +class SinglePersonSyntheticGridRule: + name: str + template_min_max_age: int + template_max_max_age: int + target_ages: tuple[int, ...] + ss_quantiles: tuple[float, ...] + payroll_quantiles: tuple[float, ...] + template_ss_state: ConstraintState = "positive" + template_payroll_state: ConstraintState = "any" + payroll_donor_min_max_age: int = 45 + payroll_donor_max_max_age: int = 64 + payroll_donor_ss_state: ConstraintState = "nonpositive" + payroll_donor_payroll_state: ConstraintState = "positive" + payroll_scale_factors: tuple[float, ...] = (1.0,) + clone_weight_scale: float = 0.1 + + @dataclass(frozen=True) class SupportAugmentationProfile: name: str description: str - rules: tuple[AgeShiftCloneRule | CompositePayrollRule, ...] + rules: tuple[ + AgeShiftCloneRule | CompositePayrollRule | SinglePersonSyntheticGridRule, ... + ] LATE_CLONE_V1 = SupportAugmentationProfile( @@ -259,12 +279,88 @@ class SupportAugmentationProfile: ), ) +LATE_SYNTHETIC_GRID_V1 = SupportAugmentationProfile( + name="late-synthetic-grid-v1", + description=( + "Appended synthetic single-person older households on an explicit " + "age/SS/payroll grid, preserving the base CPS support untouched." + ), + rules=( + SinglePersonSyntheticGridRule( + name="single_75_84_grid", + template_min_max_age=75, + template_max_max_age=84, + target_ages=(77, 82), + ss_quantiles=(0.25, 0.5, 0.75), + payroll_quantiles=(0.25, 0.5, 0.75), + template_ss_state="positive", + template_payroll_state="any", + payroll_donor_min_max_age=45, + payroll_donor_max_max_age=64, + clone_weight_scale=0.1, + ), + SinglePersonSyntheticGridRule( + name="single_85_plus_grid", + template_min_max_age=85, + template_max_max_age=85, + target_ages=(85,), + ss_quantiles=(0.25, 0.5, 0.75), + payroll_quantiles=(0.25, 0.5, 0.75), + template_ss_state="positive", + template_payroll_state="any", + payroll_donor_min_max_age=45, + payroll_donor_max_max_age=64, + clone_weight_scale=0.08, + ), + ), +) + +LATE_SYNTHETIC_GRID_V2 = SupportAugmentationProfile( + name="late-synthetic-grid-v2", + description=( + "Appended synthetic older-worker grid with more extreme payroll levels " + "to test whether older payroll intensity alone can move the late frontier." + ), + rules=( + SinglePersonSyntheticGridRule( + name="single_75_84_extreme_grid", + template_min_max_age=75, + template_max_max_age=84, + target_ages=(77, 82), + ss_quantiles=(0.25, 0.5, 0.75), + payroll_quantiles=(0.25, 0.5, 0.75), + payroll_scale_factors=(1.0, 2.0, 4.0), + template_ss_state="positive", + template_payroll_state="any", + payroll_donor_min_max_age=45, + payroll_donor_max_max_age=64, + clone_weight_scale=0.08, + ), + SinglePersonSyntheticGridRule( + name="single_85_plus_extreme_grid", + template_min_max_age=85, + template_max_max_age=85, + target_ages=(85,), + ss_quantiles=(0.25, 0.5, 0.75), + payroll_quantiles=(0.25, 0.5, 0.75), + payroll_scale_factors=(1.0, 2.0, 4.0), + template_ss_state="positive", + template_payroll_state="any", + payroll_donor_min_max_age=45, + payroll_donor_max_max_age=64, + clone_weight_scale=0.06, + ), + ), +) + NAMED_SUPPORT_AUGMENTATION_PROFILES = { LATE_CLONE_V1.name: LATE_CLONE_V1, LATE_CLONE_V2.name: LATE_CLONE_V2, LATE_COMPOSITE_V1.name: LATE_COMPOSITE_V1, LATE_COMPOSITE_V2.name: LATE_COMPOSITE_V2, + LATE_SYNTHETIC_GRID_V1.name: LATE_SYNTHETIC_GRID_V1, + LATE_SYNTHETIC_GRID_V2.name: LATE_SYNTHETIC_GRID_V2, } @@ -330,6 +426,11 @@ def household_support_summary( summary["payroll_total"] = summary[ [_period_column(component, base_year) for component in PAYROLL_COMPONENTS] ].sum(axis=1) + summary["household_size"] = ( + input_df.groupby(household_id_col, sort=False) + [_period_column(PERSON_ID_COLUMN, base_year)] + .count() + ) return summary @@ -343,11 +444,26 @@ def _match_state(values: pd.Series, state: ConstraintState) -> pd.Series: raise ValueError(f"Unsupported state '{state}'") +def _age_range_mask( + ages: pd.Series, + *, + min_max_age: int, + max_max_age: int, +) -> pd.Series: + if min_max_age == 85 and max_max_age == 85: + return ages >= 85 + return ages.between(min_max_age, max_max_age) + + def select_donor_households( summary: pd.DataFrame, rule: AgeShiftCloneRule, ) -> pd.Index: - age_mask = summary["max_age"].between(rule.min_max_age, rule.max_max_age) + age_mask = _age_range_mask( + summary["max_age"], + min_max_age=rule.min_max_age, + max_max_age=rule.max_max_age, + ) positive_weight_mask = summary["baseline_weight"] > 0 ss_mask = _match_state(summary["ss_total"], rule.ss_state) payroll_mask = _match_state(summary["payroll_total"], rule.payroll_state) @@ -362,7 +478,11 @@ def select_households_for_composite_rule( ss_state: ConstraintState, payroll_state: ConstraintState, ) -> pd.Index: - age_mask = summary["max_age"].between(min_max_age, max_max_age) + age_mask = _age_range_mask( + summary["max_age"], + min_max_age=min_max_age, + max_max_age=max_max_age, + ) positive_weight_mask = summary["baseline_weight"] > 0 ss_mask = _match_state(summary["ss_total"], ss_state) payroll_mask = _match_state(summary["payroll_total"], payroll_state) @@ -643,6 +763,141 @@ def synthesize_composite_households( ) +def synthesize_single_person_grid_households( + input_df: pd.DataFrame, + *, + base_year: int, + summary: pd.DataFrame, + rule: SinglePersonSyntheticGridRule, + id_counters: dict[str, int] | None = None, +) -> tuple[pd.DataFrame, dict[str, int], dict[str, Any]]: + template_ids = select_households_for_composite_rule( + summary[summary["household_size"] == 1], + min_max_age=rule.template_min_max_age, + max_max_age=rule.template_max_max_age, + ss_state=rule.template_ss_state, + payroll_state=rule.template_payroll_state, + ) + donor_ids = select_households_for_composite_rule( + summary[summary["household_size"] == 1], + min_max_age=rule.payroll_donor_min_max_age, + max_max_age=rule.payroll_donor_max_max_age, + ss_state=rule.payroll_donor_ss_state, + payroll_state=rule.payroll_donor_payroll_state, + ) + if template_ids.empty or donor_ids.empty: + return ( + input_df.iloc[0:0].copy(), + id_counters.copy() if id_counters is not None else {}, + { + "rule": rule.name, + "template_household_count": int(len(template_ids)), + "payroll_donor_household_count": int(len(donor_ids)), + "synthetic_household_count": 0, + "synthetic_person_count": 0, + }, + ) + + household_id_col = _period_column("household_id", base_year) + age_col = _period_column("age", base_year) + person_id_col = _period_column(PERSON_ID_COLUMN, base_year) + household_weight_col = _period_column("household_weight", base_year) + person_weight_col = _period_column("person_weight", base_year) + employment_col = _period_column("employment_income_before_lsr", base_year) + self_employment_col = _period_column("self_employment_income_before_lsr", base_year) + qbi_col = _period_column("w2_wages_from_qualified_business", base_year) + + template_rows = ( + input_df[input_df[household_id_col].isin(template_ids)] + .sort_values([age_col, household_id_col]) + .copy() + ) + ss_values = summary.loc[template_ids, "ss_total"].astype(float) + payroll_values = summary.loc[donor_ids, "payroll_total"].astype(float) + ss_targets = [float(ss_values.quantile(q)) for q in rule.ss_quantiles] + payroll_targets = [ + float(payroll_values.quantile(q) * scale_factor) + for q in rule.payroll_quantiles + for scale_factor in rule.payroll_scale_factors + ] + + next_ids = ( + id_counters.copy() + if id_counters is not None + else { + entity_name: _next_entity_id( + input_df[_period_column(columns[0], base_year)] + ) + for entity_name, columns in ENTITY_ID_COLUMNS.items() + } + ) + if "person" not in next_ids: + next_ids["person"] = _next_entity_id(input_df[person_id_col]) + + synthetic_rows = [] + template_records = template_rows.to_dict("records") + template_index = 0 + for target_age in rule.target_ages: + for ss_target in ss_targets: + for payroll_target in payroll_targets: + base_row = template_records[template_index % len(template_records)].copy() + template_index += 1 + household_id = next_ids["household"] + next_ids["household"] += 1 + person_id = next_ids["person"] + next_ids["person"] += 1 + + for entity_name, columns in ENTITY_ID_COLUMNS.items(): + entity_id = next_ids[entity_name] + next_ids[entity_name] += 1 + for raw_column in columns: + column = _period_column(raw_column, base_year) + base_row[column] = entity_id + + base_row[household_id_col] = household_id + base_row[_period_column("person_household_id", base_year)] = household_id + base_row[person_id_col] = person_id + base_row[age_col] = float(target_age) + base_row[_period_column("social_security_retirement", base_year)] = float( + ss_target + ) + base_row[_period_column("social_security_disability", base_year)] = 0.0 + base_row[_period_column("social_security_survivors", base_year)] = 0.0 + base_row[_period_column("social_security_dependents", base_year)] = 0.0 + base_row[employment_col] = float(payroll_target) + base_row[self_employment_col] = 0.0 + if qbi_col in base_row: + base_row[qbi_col] = 0.0 + if household_weight_col in base_row: + base_row[household_weight_col] = ( + float(base_row[household_weight_col]) * rule.clone_weight_scale + ) + if person_weight_col in base_row: + base_row[person_weight_col] = ( + float(base_row[person_weight_col]) * rule.clone_weight_scale + ) + synthetic_rows.append(base_row) + + synthetic_df = pd.DataFrame(synthetic_rows, columns=input_df.columns) + return ( + synthetic_df, + next_ids, + { + "rule": rule.name, + "template_household_count": int(len(template_ids)), + "payroll_donor_household_count": int(len(donor_ids)), + "synthetic_household_count": int( + synthetic_df[household_id_col].nunique() + ), + "synthetic_person_count": int(len(synthetic_df)), + "target_age_count": int(len(rule.target_ages)), + "ss_grid_size": int(len(rule.ss_quantiles)), + "payroll_grid_size": int(len(rule.payroll_quantiles)), + "payroll_scale_factor_count": int(len(rule.payroll_scale_factors)), + }, + ) + + def augment_input_dataframe( input_df: pd.DataFrame, *, @@ -693,15 +948,27 @@ def augment_input_dataframe( ) continue - composite_df, id_counters, composite_report = synthesize_composite_households( + if isinstance(rule, CompositePayrollRule): + composite_df, id_counters, composite_report = synthesize_composite_households( + input_df, + base_year=base_year, + summary=summary, + rule=rule, + id_counters=id_counters, + ) + clone_frames.append(composite_df) + rule_reports.append(composite_report) + continue + + synthetic_df, id_counters, synthetic_report = synthesize_single_person_grid_households( input_df, base_year=base_year, summary=summary, rule=rule, id_counters=id_counters, ) - clone_frames.append(composite_df) - rule_reports.append(composite_report) + clone_frames.append(synthetic_df) + rule_reports.append(synthetic_report) if clone_frames: augmented_df = pd.concat([input_df, *clone_frames], ignore_index=True) diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index 1f988c1fb..ec83f377d 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -36,6 +36,7 @@ from policyengine_us_data.datasets.cps.long_term.support_augmentation import ( AgeShiftCloneRule, CompositePayrollRule, + SinglePersonSyntheticGridRule, SupportAugmentationProfile, augment_input_dataframe, household_support_summary, @@ -182,6 +183,68 @@ def test_support_augmentation_synthesizes_composite_payroll_household(): assert cloned_rows["employment_income_before_lsr__2024"].sum() == pytest.approx(37_000.0) +def test_support_augmentation_appends_single_person_synthetic_grid_households(): + import pandas as pd + + df = pd.DataFrame( + { + "person_id__2024": [101, 201, 301], + "household_id__2024": [1, 2, 3], + "person_household_id__2024": [1, 2, 3], + "family_id__2024": [11.0, 21.0, 31.0], + "person_family_id__2024": [11, 21, 31], + "tax_unit_id__2024": [101, 201, 301], + "person_tax_unit_id__2024": [101, 201, 301], + "spm_unit_id__2024": [1001, 2001, 3001], + "person_spm_unit_id__2024": [1001, 2001, 3001], + "marital_unit_id__2024": [501, 601, 701], + "person_marital_unit_id__2024": [501, 601, 701], + "age__2024": [78.0, 86.0, 60.0], + "household_weight__2024": [10.0, 8.0, 5.0], + "person_weight__2024": [10.0, 8.0, 5.0], + "social_security_retirement__2024": [20_000.0, 24_000.0, 0.0], + "social_security_disability__2024": [0.0, 0.0, 0.0], + "social_security_survivors__2024": [0.0, 0.0, 0.0], + "social_security_dependents__2024": [0.0, 0.0, 0.0], + "employment_income_before_lsr__2024": [0.0, 0.0, 50_000.0], + "self_employment_income_before_lsr__2024": [0.0, 0.0, 0.0], + "w2_wages_from_qualified_business__2024": [0.0, 0.0, 0.0], + } + ) + profile = SupportAugmentationProfile( + name="grid-profile", + description="Toy single-person synthetic grid.", + rules=( + SinglePersonSyntheticGridRule( + name="older_grid", + template_min_max_age=75, + template_max_max_age=86, + target_ages=(77, 85), + ss_quantiles=(0.5,), + payroll_quantiles=(0.5,), + template_ss_state="positive", + template_payroll_state="any", + payroll_donor_min_max_age=55, + payroll_donor_max_max_age=64, + clone_weight_scale=0.2, + ), + ), + ) + augmented_df, report = augment_input_dataframe( + df, + base_year=2024, + profile=profile, + ) + assert report["base_household_count"] == 3 + assert report["augmented_household_count"] == 5 + synthetic_household_ids = set(augmented_df["household_id__2024"].unique()) - {1, 2, 3} + assert len(synthetic_household_ids) == 2 + synthetic_rows = augmented_df[augmented_df["household_id__2024"].isin(synthetic_household_ids)] + assert set(synthetic_rows["age__2024"].tolist()) == {77.0, 85.0} + assert set(synthetic_rows["social_security_retirement__2024"].tolist()) == {22_000.0} + assert set(synthetic_rows["employment_income_before_lsr__2024"].tolist()) == {50_000.0} + + def test_age_bin_helpers_preserve_population_totals(): bins = build_age_bins(n_ages=86, bucket_size=5) assert bins[0] == (0, 5) From c99ccbac8b4e0e5c858a4d580234c8853a75f642 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 31 Mar 2026 14:10:47 -0400 Subject: [PATCH 13/31] Benchmark long-run TOB outside calibration --- .../cps/long_term/ASSUMPTION_COMPARISON.md | 26 +++++++++ .../datasets/cps/long_term/README.md | 15 +++-- .../long_term/assess_calibration_frontier.py | 2 +- .../cps/long_term/calibration_profiles.py | 15 +++-- .../evaluate_support_augmentation.py | 2 +- .../cps/long_term/run_household_projection.py | 57 +++++++++++++++---- .../test_long_term_calibration_contract.py | 45 +++++++++------ 7 files changed, 121 insertions(+), 41 deletions(-) create mode 100644 policyengine_us_data/datasets/cps/long_term/ASSUMPTION_COMPARISON.md diff --git a/policyengine_us_data/datasets/cps/long_term/ASSUMPTION_COMPARISON.md b/policyengine_us_data/datasets/cps/long_term/ASSUMPTION_COMPARISON.md new file mode 100644 index 000000000..4a9764488 --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/ASSUMPTION_COMPARISON.md @@ -0,0 +1,26 @@ +# Long-Run Calibration Assumption Comparison + +This note distinguishes between: + +- hard microsimulation calibration targets, which directly shape household weights +- post-calibration benchmarks, which are compared against published source packages without being forced into the weights + +The split is intentional. Long-run Taxation of Benefits (TOB) depends on tax-side assumptions that are not fully published as a household-level rule schedule in the Trustees/OACT materials. + +| Component | Current `policyengine-us-data` approach | Trustees / OACT published approach | Calibration use | +| --- | --- | --- | --- | +| Population by age | SSA single-year age projections | SSA single-year age projections | Hard target | +| OASDI benefits | Named long-term target source package | Trustees or OACT-patched annual OASDI path | Hard target | +| Taxable payroll | Named long-term target source package | Trustees annual taxable payroll path | Hard target | +| Social Security benefit-tax thresholds | Literal current-law statutory thresholds remain fixed in nominal dollars | Trustees also describe statutory thresholds as remaining fixed in nominal dollars | Not separately targeted | +| Federal income-tax brackets | PolicyEngine tax simulation using its own income-tax parameter uprating path | Trustees assume periodic future bracket adjustments; after the tenth year, tax brackets are assumed to rise with average wages to avoid indefinite bracket creep | Not hard-targeted | +| OASDI TOB | Computed from the tax microsimulation and compared to the selected target source | Trustees/OACT publish annual revenue paths or ratios, but not a full public household-level micro rule schedule | Post-calibration benchmark | +| HI TOB | Computed from the tax microsimulation and compared to the selected target source | Trustees publish current-law HI TOB path; OACT OBBBA updates do not currently provide a full public annual HI replacement series | Post-calibration benchmark | +| OBBBA OASDI update | Available through named target source `oact_2025_08_05_provisional` | August 5, 2025 OACT letter provides annual OASDI changes through 2100 | Benchmark / target-source input | +| OBBBA HI update | Provisional bridge only in named target source | No equivalent full public annual HI replacement path located yet | Benchmark only | + +## Practical interpretation + +- `ss-payroll` is the core hard-target profile. +- `ss-payroll-tob` now means: calibrate on age + OASDI benefits + taxable payroll, then benchmark TOB against the selected source package. +- TOB remains important for model comparison, but it is no longer treated as a weight-identifying target in the long-run microsimulation contract. diff --git a/policyengine_us_data/datasets/cps/long_term/README.md b/policyengine_us_data/datasets/cps/long_term/README.md index e8de8dfe0..11f082824 100644 --- a/policyengine_us_data/datasets/cps/long_term/README.md +++ b/policyengine_us_data/datasets/cps/long_term/README.md @@ -6,7 +6,7 @@ Run projections using `run_household_projection.py`: ```bash -# Recommended: named profile with all constraint types +# Recommended: named profile with TOB benchmarked post-calibration python run_household_projection.py 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --save-h5 # IPF with only age distribution constraints (faster, less accurate) @@ -31,8 +31,8 @@ python run_household_projection.py 2100 --profile ss - `age-only`: IPF age-only calibration - `ss`: positive entropy calibration with age + Social Security - `ss-payroll`: positive entropy calibration with age + Social Security + taxable payroll -- `ss-payroll-tob`: positive entropy calibration with age + Social Security + taxable payroll + TOB -- `ss-payroll-tob-h6`: positive entropy calibration with age + Social Security + taxable payroll + TOB + H6 +- `ss-payroll-tob`: positive entropy calibration with age + Social Security + taxable payroll, with TOB benchmarked after calibration +- `ss-payroll-tob-h6`: positive entropy calibration with age + Social Security + taxable payroll + H6, with TOB benchmarked after calibration **Validation contract:** - Economic-targeted profiles no longer silently pretend an IPF fallback is equivalent to GREG. @@ -56,7 +56,7 @@ python run_household_projection.py 2100 --profile ss **Positive Entropy Calibration** - Solves for strictly positive weights matching multiple constraints simultaneously -- Can enforce age distribution + Social Security benefits + taxable payroll + TOB revenue +- Can enforce age distribution + Social Security benefits + taxable payroll - Uses dual optimization to minimize divergence from baseline weights - **Recommended** for publishable long-term projections @@ -82,11 +82,12 @@ python run_household_projection.py 2100 --profile ss - Calculated as: `taxable_earnings_for_social_security` + `social_security_taxable_self_employment_income` - Source: selected long-term target source package -4. **TOB Revenue** (`--use-tob`, GREG only) +4. **TOB Revenue** (`--use-tob`, legacy hard-target mode only) - Taxation of Benefits revenue for OASDI and Medicare HI trust funds - OASDI: `tob_revenue_oasdi` (tier 1 taxation, 0-50% of benefits) - HI: `tob_revenue_medicare_hi` (tier 2 taxation, 50-85% of benefits) - Source: selected long-term target source package + - Recommended usage: benchmark after calibration rather than use as a hard weight target --- @@ -106,6 +107,7 @@ python run_household_projection.py 2100 --profile ss - `SSPopJul_TR2024.csv` - Population projections 2025-2100 by single year of age - `long_term_target_sources/trustees_2025_current_law.csv` - explicit frozen Trustees/current-law package - `long_term_target_sources/sources.json` - provenance metadata for named source packages +- `ASSUMPTION_COMPARISON.md` - side-by-side summary of our calibration assumptions versus Trustees/OACT --- @@ -141,7 +143,8 @@ For each projection year (2025-2100): 2. **Uprate variables** - PolicyEngine automatically uprates income, thresholds, etc. to target year 3. **Calculate values** - Income tax, Social Security, taxable payroll at household level 4. **Calibrate weights** - Adjust household weights to match SSA demographic/economic targets -5. **Aggregate results** - Apply calibrated weights to calculate national totals +5. **Benchmark TOB** - Compare modeled OASDI/HI TOB to the selected target source without forcing it into the weights +6. **Aggregate results** - Apply calibrated weights to calculate national totals **Key innovation:** Household-level calculations avoid person→household aggregation issues, maintaining consistency across all variables. diff --git a/policyengine_us_data/datasets/cps/long_term/assess_calibration_frontier.py b/policyengine_us_data/datasets/cps/long_term/assess_calibration_frontier.py index 63e6ac627..04f024850 100644 --- a/policyengine_us_data/datasets/cps/long_term/assess_calibration_frontier.py +++ b/policyengine_us_data/datasets/cps/long_term/assess_calibration_frontier.py @@ -42,7 +42,7 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument( "--profile", - default="ss-payroll-tob", + default="ss-payroll", help="Named calibration profile to assess.", ) parser.add_argument( diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py index 7a5d45ea2..59f2c27d6 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py @@ -35,6 +35,7 @@ class CalibrationProfile: use_payroll: bool use_h6_reform: bool use_tob: bool + benchmark_tob: bool = False allow_greg_fallback: bool = False max_constraint_error_pct: float = 0.1 max_age_error_pct: float = 0.1 @@ -113,6 +114,7 @@ def to_dict(self) -> dict[str, Any]: use_payroll=False, use_h6_reform=False, use_tob=False, + benchmark_tob=False, allow_greg_fallback=False, min_positive_household_count=1000, min_effective_sample_size=75.0, @@ -128,6 +130,7 @@ def to_dict(self) -> dict[str, Any]: use_payroll=False, use_h6_reform=False, use_tob=False, + benchmark_tob=False, max_negative_weight_pct=0.0, min_positive_household_count=1000, min_effective_sample_size=75.0, @@ -144,6 +147,7 @@ def to_dict(self) -> dict[str, Any]: use_payroll=True, use_h6_reform=False, use_tob=False, + benchmark_tob=False, max_negative_weight_pct=0.0, min_positive_household_count=1000, min_effective_sample_size=75.0, @@ -153,13 +157,14 @@ def to_dict(self) -> dict[str, Any]: ), "ss-payroll-tob": CalibrationProfile( name="ss-payroll-tob", - description="Age, Social Security, taxable payroll, and TOB using positive entropy calibration.", + description="Age, Social Security, and taxable payroll using positive entropy calibration, with TOB benchmarked post-calibration against the selected long-term target source.", calibration_method="entropy", use_greg=False, use_ss=True, use_payroll=True, use_h6_reform=False, - use_tob=True, + use_tob=False, + benchmark_tob=True, max_negative_weight_pct=0.0, min_positive_household_count=1000, min_effective_sample_size=75.0, @@ -169,13 +174,14 @@ def to_dict(self) -> dict[str, Any]: ), "ss-payroll-tob-h6": CalibrationProfile( name="ss-payroll-tob-h6", - description="Age, Social Security, taxable payroll, TOB, and H6 using positive entropy calibration.", + description="Age, Social Security, taxable payroll, and H6 using positive entropy calibration, with TOB benchmarked post-calibration against the selected long-term target source.", calibration_method="entropy", use_greg=False, use_ss=True, use_payroll=True, use_h6_reform=True, - use_tob=True, + use_tob=False, + benchmark_tob=True, max_negative_weight_pct=0.0, min_positive_household_count=1000, min_effective_sample_size=75.0, @@ -275,6 +281,7 @@ def build_profile_from_flags( use_payroll=use_payroll, use_h6_reform=use_h6_reform, use_tob=use_tob, + benchmark_tob=False, ) diff --git a/policyengine_us_data/datasets/cps/long_term/evaluate_support_augmentation.py b/policyengine_us_data/datasets/cps/long_term/evaluate_support_augmentation.py index de4ed2860..b70624378 100644 --- a/policyengine_us_data/datasets/cps/long_term/evaluate_support_augmentation.py +++ b/policyengine_us_data/datasets/cps/long_term/evaluate_support_augmentation.py @@ -148,7 +148,7 @@ def main() -> None: parser.add_argument("year", type=int, help="Projection year to evaluate.") parser.add_argument( "--profile", - default="ss-payroll-tob", + default="ss-payroll", help="Calibration profile to evaluate.", ) parser.add_argument( diff --git a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py index 8983a947b..c1eeb60a7 100644 --- a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py +++ b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py @@ -15,7 +15,7 @@ --use-ss: Include Social Security benefit totals as calibration target (requires --greg) --use-payroll: Include taxable payroll totals as calibration target (requires --greg) --use-h6-reform: Include H6 reform income impact ratio as calibration target (requires --greg) - --use-tob: Include TOB (Taxation of Benefits) revenue as calibration target (requires --greg) + --use-tob: Include TOB (Taxation of Benefits) revenue as a hard calibration target (requires --greg) --save-h5: Save year-specific .h5 files with calibrated weights to ./projected_datasets/ Examples: @@ -335,11 +335,12 @@ def create_h6_reform(): USE_PAYROLL = PROFILE.use_payroll USE_H6_REFORM = PROFILE.use_h6_reform USE_TOB = PROFILE.use_tob +BENCHMARK_TOB = PROFILE.benchmark_tob if USE_H6_REFORM: from ssa_data import load_h6_income_rate_change -if USE_TOB: +if USE_TOB or BENCHMARK_TOB: from ssa_data import load_hi_tob_projections, load_oasdi_tob_projections if USE_GREG: @@ -373,6 +374,8 @@ def create_h6_reform(): print(f" Including H6 reform income impact constraint: Yes") if USE_TOB: print(f" Including TOB revenue constraint: Yes") +elif BENCHMARK_TOB: + print(f" Benchmarking TOB after calibration: Yes") if SAVE_H5: print(f" Saving year-specific .h5 files: Yes (to {OUTPUT_DIR}/)") os.makedirs(OUTPUT_DIR, exist_ok=True) @@ -549,7 +552,7 @@ def create_h6_reform(): oasdi_tob_target = None hi_tob_values = None hi_tob_target = None - if USE_TOB: + if USE_TOB or BENCHMARK_TOB: oasdi_tob_hh = sim.calculate( "tob_revenue_oasdi", period=year, map_to="household" ) @@ -604,10 +607,10 @@ def create_h6_reform(): payroll_target=payroll_target, h6_income_values=h6_income_values, h6_revenue_target=h6_revenue_target, - oasdi_tob_values=oasdi_tob_values, - oasdi_tob_target=oasdi_tob_target, - hi_tob_values=hi_tob_values, - hi_tob_target=hi_tob_target, + oasdi_tob_values=oasdi_tob_values if USE_TOB else None, + oasdi_tob_target=oasdi_tob_target if USE_TOB else None, + hi_tob_values=hi_tob_values if USE_TOB else None, + hi_tob_target=hi_tob_target if USE_TOB else None, n_ages=X_current.shape[1], max_iters=100, tol=1e-6, @@ -633,11 +636,30 @@ def create_h6_reform(): payroll_target=payroll_target, h6_income_values=h6_income_values, h6_revenue_target=h6_revenue_target, - oasdi_tob_values=oasdi_tob_values, - oasdi_tob_target=oasdi_tob_target, - hi_tob_values=hi_tob_values, - hi_tob_target=hi_tob_target, + oasdi_tob_values=oasdi_tob_values if USE_TOB else None, + oasdi_tob_target=oasdi_tob_target if USE_TOB else None, + hi_tob_values=hi_tob_values if USE_TOB else None, + hi_tob_target=hi_tob_target if USE_TOB else None, ) + if BENCHMARK_TOB and oasdi_tob_values is not None and hi_tob_values is not None: + calibration_audit["benchmarks"] = { + "oasdi_tob": { + "target": float(oasdi_tob_target), + "achieved": float(np.sum(oasdi_tob_values * w_new)), + }, + "hi_tob": { + "target": float(hi_tob_target), + "achieved": float(np.sum(hi_tob_values * w_new)), + }, + } + for benchmark in calibration_audit["benchmarks"].values(): + benchmark["error"] = benchmark["achieved"] - benchmark["target"] + benchmark["pct_error"] = ( + 0.0 + if benchmark["target"] == 0 + else (benchmark["error"] / benchmark["target"] * 100) + ) + benchmark["source"] = TARGET_SOURCE calibration_audit["calibration_quality"] = classify_calibration_quality( calibration_audit, PROFILE, @@ -702,6 +724,19 @@ def create_h6_reform(): f"(error: ${abs(hi_stats['error']) / 1e6:.1f}M, " f"{hi_stats['pct_error']:.3f}%)" ) + if year in display_years and BENCHMARK_TOB: + oasdi_stats = calibration_audit["benchmarks"]["oasdi_tob"] + hi_stats = calibration_audit["benchmarks"]["hi_tob"] + print( + f" [DEBUG {year}] OASDI TOB benchmark: ${oasdi_stats['achieved'] / 1e9:.1f}B " + f"(gap: ${abs(oasdi_stats['error']) / 1e6:.1f}M, " + f"{oasdi_stats['pct_error']:.3f}%)" + ) + print( + f" [DEBUG {year}] HI TOB benchmark: ${hi_stats['achieved'] / 1e9:.1f}B " + f"(gap: ${abs(hi_stats['error']) / 1e6:.1f}M, " + f"{hi_stats['pct_error']:.3f}%)" + ) weights_matrix[:, year_idx] = w_new baseline_weights_matrix[:, year_idx] = baseline_weights diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index ec83f377d..6174af33b 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -84,7 +84,8 @@ def test_named_profile_lookup(): assert profile.use_greg is False assert profile.use_ss is True assert profile.use_payroll is True - assert profile.use_tob is True + assert profile.use_tob is False + assert profile.benchmark_tob is True assert profile.use_h6_reform is False assert profile.max_negative_weight_pct == 0.0 assert profile.approximate_windows[0].age_bucket_size == 5 @@ -275,8 +276,8 @@ def test_legacy_flags_map_to_named_profile(): use_h6_reform=False, use_tob=True, ) - assert profile.name == "ss-payroll-tob" - assert profile.calibration_method == "entropy" + assert profile.name == "custom-ss-payroll-tob" + assert profile.calibration_method == "ipf" def test_strict_greg_failure_raises(): @@ -356,9 +357,7 @@ def test_classify_calibration_quality_marks_invalid_audit_approximate(): "negative_weight_pct": 0.0, "constraints": { "ss_total": {"pct_error": 0.0}, - "payroll_total": {"pct_error": 0.0}, - "oasdi_tob": {"pct_error": 0.5}, - "hi_tob": {"pct_error": 0.0}, + "payroll_total": {"pct_error": 0.5}, }, }, profile, @@ -377,8 +376,6 @@ def test_entropy_profile_rejects_negative_weights(): "constraints": { "ss_total": {"pct_error": 0.0}, "payroll_total": {"pct_error": 0.0}, - "oasdi_tob": {"pct_error": 0.0}, - "hi_tob": {"pct_error": 0.0}, }, }, profile, @@ -400,8 +397,6 @@ def test_support_thresholds_reject_concentrated_weights(): "constraints": { "ss_total": {"pct_error": 0.0}, "payroll_total": {"pct_error": 0.0}, - "oasdi_tob": {"pct_error": 0.0}, - "hi_tob": {"pct_error": 0.0}, }, }, profile, @@ -427,8 +422,6 @@ def test_classify_calibration_quality_marks_support_collapse_aggregate(): "constraints": { "ss_total": {"pct_error": 0.0}, "payroll_total": {"pct_error": 0.0}, - "oasdi_tob": {"pct_error": 0.0}, - "hi_tob": {"pct_error": 0.0}, }, }, profile, @@ -447,8 +440,6 @@ def test_approximate_window_is_year_bounded(): "constraints": { "ss_total": {"pct_error": 0.0}, "payroll_total": {"pct_error": 3.0}, - "oasdi_tob": {"pct_error": 3.0}, - "hi_tob": {"pct_error": 3.0}, }, }, profile, @@ -464,8 +455,6 @@ def test_approximate_window_is_year_bounded(): "constraints": { "ss_total": {"pct_error": 0.0}, "payroll_total": {"pct_error": 3.0}, - "oasdi_tob": {"pct_error": 3.0}, - "hi_tob": {"pct_error": 3.0}, }, }, profile, @@ -490,8 +479,6 @@ def test_normalize_metadata_harmonizes_lp_fallback_labels(): "constraints": { "ss_total": {"pct_error": 0.368}, "payroll_total": {"pct_error": 0.0}, - "oasdi_tob": {"pct_error": 0.0}, - "hi_tob": {"pct_error": 0.0}, }, }, } @@ -572,6 +559,28 @@ def test_manifest_updates_and_rejects_profile_mismatch(tmp_path): assert rebuilt["years"] == [2026, 2027] +def test_benchmarked_tob_does_not_affect_quality_classification(): + profile = get_profile("ss-payroll-tob") + quality = classify_calibration_quality( + { + "fell_back_to_ipf": False, + "age_max_pct_error": 0.0, + "negative_weight_pct": 0.0, + "constraints": { + "ss_total": {"pct_error": 0.0}, + "payroll_total": {"pct_error": 0.0}, + }, + "benchmarks": { + "oasdi_tob": {"pct_error": 12.0}, + "hi_tob": {"pct_error": -9.0}, + }, + }, + profile, + year=2035, + ) + assert quality == "exact" + + def test_entropy_calibration_produces_nonnegative_weights_and_hits_targets(): X = np.array( [ From 642752cb17d758abdb0752e1ea45698a0ee2e97a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 31 Mar 2026 15:12:26 -0400 Subject: [PATCH 14/31] Try denser late-tail approximate calibration --- .../datasets/cps/long_term/calibration.py | 303 +++++++++++++++++- .../test_long_term_calibration_contract.py | 31 +- 2 files changed, 325 insertions(+), 9 deletions(-) diff --git a/policyengine_us_data/datasets/cps/long_term/calibration.py b/policyengine_us_data/datasets/cps/long_term/calibration.py index 7dded31cd..7cb030f8c 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration.py @@ -383,7 +383,223 @@ def infeasibility_error(prefix): return weights, iterations -def calibrate_entropy_approximate( +def calibrate_entropy_bounded( + X, + y_target, + baseline_weights, + ss_values=None, + ss_target=None, + payroll_values=None, + payroll_target=None, + h6_income_values=None, + h6_revenue_target=None, + oasdi_tob_values=None, + oasdi_tob_target=None, + hi_tob_values=None, + hi_tob_target=None, + n_ages=86, + max_constraint_error_pct=0.0, + max_iters=500, + tol=1e-9, + warm_weights=None, +): + """ + Approximate positive calibration via entropy balancing inside an error box. + + This keeps the entropy objective, but relaxes the constraints to + `|A'w - y| / scale <= epsilon` where `epsilon` is the allowed maximum + percent error and `scale = max(|target|, 1)`. It is much denser than the + LP minimax solution and is therefore a better approximate microsimulation + fallback when exact calibration is infeasible. + """ + + aux_df, controls = _build_constraint_dataframe_and_controls( + X, + y_target, + ss_values=ss_values, + ss_target=ss_target, + payroll_values=payroll_values, + payroll_target=payroll_target, + h6_income_values=h6_income_values, + h6_revenue_target=h6_revenue_target, + oasdi_tob_values=oasdi_tob_values, + oasdi_tob_target=oasdi_tob_target, + hi_tob_values=hi_tob_values, + hi_tob_target=hi_tob_target, + n_ages=n_ages, + ) + + A = aux_df.to_numpy(dtype=float) + targets = np.array(list(controls.values()), dtype=float) + baseline_weights = np.asarray(baseline_weights, dtype=float) + + epsilon = max(float(max_constraint_error_pct) / 100.0, 0.0) + scales = np.maximum(np.abs(targets), 1.0) + A_scaled = A / scales + targets_scaled = targets / scales + lower_bounds = targets_scaled - epsilon + upper_bounds = targets_scaled + epsilon + + gram = A_scaled.T @ (baseline_weights[:, None] * A_scaled) + gram += np.eye(gram.shape[0]) * 1e-12 + beta0 = np.linalg.solve( + gram, + targets_scaled - (A_scaled.T @ baseline_weights), + ) + + def objective_and_gradient(z): + n_constraints = len(targets_scaled) + alpha = z[:n_constraints] + gamma = z[n_constraints:] + beta = gamma - alpha + eta = np.clip(A_scaled @ beta, -700, 700) + exp_eta = np.exp(eta) + weights = baseline_weights * exp_eta + achieved = A_scaled.T @ weights + objective = float(np.sum(weights) + upper_bounds @ alpha - lower_bounds @ gamma) + gradient = np.concatenate( + [ + upper_bounds - achieved, + achieved - lower_bounds, + ] + ) + return objective, gradient, weights + + starts = [ + np.zeros(len(targets_scaled) * 2, dtype=float), + np.concatenate( + [ + np.maximum(-beta0, 0.0), + np.maximum(beta0, 0.0), + ] + ), + ] + + for weights_start in warm_weights or []: + weights_start = np.asarray(weights_start, dtype=float) + if weights_start.shape != baseline_weights.shape: + continue + ratios = np.clip( + weights_start / np.maximum(baseline_weights, 1e-300), + 1e-300, + 1e300, + ) + beta_start, *_ = np.linalg.lstsq( + A_scaled, + np.log(ratios), + rcond=None, + ) + starts.append( + np.concatenate( + [ + np.maximum(-beta_start, 0.0), + np.maximum(beta_start, 0.0), + ] + ) + ) + + best_result = None + best_weights = None + best_max_error_pct = float("inf") + + for start in starts: + result = optimize.minimize( + lambda z: objective_and_gradient(z)[0], + start, + jac=lambda z: objective_and_gradient(z)[1], + method="L-BFGS-B", + bounds=[(0.0, None)] * len(start), + options={"maxiter": max_iters, "ftol": tol}, + ) + + objective, gradient, weights = objective_and_gradient(result.x) + achieved = A_scaled.T @ weights + max_error_pct = float( + np.max(np.abs(achieved - targets_scaled)) * 100 + ) + + if max_error_pct < best_max_error_pct: + best_result = result + best_weights = weights + best_max_error_pct = max_error_pct + + if result.success and max_error_pct <= max_constraint_error_pct + 1e-6: + return ( + np.asarray(weights, dtype=float), + int(result.nit), + { + "success": True, + "best_case_max_pct_error": max_error_pct, + "status": int(result.status), + "message": result.message, + }, + ) + + if best_result is None or best_weights is None: + raise RuntimeError("Approximate bounded entropy calibration did not run.") + + raise RuntimeError( + "Approximate bounded entropy calibration failed: " + f"best achieved max relative constraint error was " + f"{best_max_error_pct:.3f}%" + ) + + +def densify_lp_solution( + A, + targets, + baseline_weights, + lp_weights, + max_constraint_error_pct, + *, + iterations=30, +): + """ + Blend an LP solution back toward the baseline while staying inside the + allowed error band. + + This preserves feasibility while avoiding the most extreme support collapse + of a pure LP basic-feasible-point solution. + """ + + A = np.asarray(A, dtype=float) + targets = np.asarray(targets, dtype=float) + baseline_weights = np.asarray(baseline_weights, dtype=float) + lp_weights = np.asarray(lp_weights, dtype=float) + + scales = np.maximum(np.abs(targets), 1.0) + A_scaled = A / scales + targets_scaled = targets / scales + + best_lambda = 0.0 + best_weights = lp_weights.copy() + best_error_pct = float( + np.max(np.abs(A_scaled.T @ best_weights - targets_scaled)) * 100 + ) + + lo = 0.0 + hi = 1.0 + for _ in range(iterations): + lam = (lo + hi) / 2.0 + candidate_weights = (1.0 - lam) * lp_weights + lam * baseline_weights + candidate_error_pct = float( + np.max(np.abs(A_scaled.T @ candidate_weights - targets_scaled)) * 100 + ) + if candidate_error_pct <= max_constraint_error_pct + 1e-6: + best_lambda = lam + best_weights = candidate_weights + best_error_pct = candidate_error_pct + lo = lam + else: + hi = lam + + return best_weights, { + "blend_lambda": best_lambda, + "best_case_max_pct_error": best_error_pct, + } + + +def calibrate_lp_minimax( X, y_target, baseline_weights, @@ -402,9 +618,9 @@ def calibrate_entropy_approximate( """ Approximate nonnegative calibration via minimax relative-error LP. - This is the robust fallback when exact positive entropy calibration is - infeasible under the current support. It returns the best nonnegative - weight vector available under the requested constraints. + This is a robust feasibility fallback and certificate, but it tends to + produce sparse extreme-point solutions. Prefer bounded entropy when we want + approximate weights to remain usable as microsimulation support. """ aux_df, controls = _build_constraint_dataframe_and_controls( @@ -553,6 +769,8 @@ def calibrate_weights( "method_used": method, "greg_attempted": method == "greg", "greg_error": None, + "entropy_error": None, + "approximate_entropy_error": None, "fell_back_to_ipf": False, "lp_fallback_used": False, "approximate_solution_used": False, @@ -619,7 +837,7 @@ def calibrate_weights( return w_new, iterations, audit except RuntimeError as error: audit["entropy_error"] = str(error) - w_new, iterations, feasibility = calibrate_entropy_approximate( + w_new, iterations, feasibility = calibrate_lp_minimax( X, y_target, baseline_weights, @@ -654,6 +872,81 @@ def calibrate_weights( f"{approximate_error_pct:.3f}% > {approximate_max_error_pct:.3f}%" ) from error + dense_lp_weights = None + dense_lp_info = None + try: + aux_df, controls = _build_constraint_dataframe_and_controls( + X, + y_target, + ss_values=ss_values, + ss_target=ss_target, + payroll_values=payroll_values, + payroll_target=payroll_target, + h6_income_values=h6_income_values, + h6_revenue_target=h6_revenue_target, + oasdi_tob_values=oasdi_tob_values, + oasdi_tob_target=oasdi_tob_target, + hi_tob_values=hi_tob_values, + hi_tob_target=hi_tob_target, + n_ages=n_ages, + ) + dense_lp_weights, dense_lp_info = densify_lp_solution( + aux_df.to_numpy(dtype=float), + np.array(list(controls.values()), dtype=float), + baseline_weights, + w_new, + approximate_max_error_pct, + ) + except Exception: + dense_lp_weights = None + dense_lp_info = None + + if approximate_max_error_pct is not None: + try: + warm_weights = [w_new] + if dense_lp_weights is not None: + warm_weights.insert(0, dense_lp_weights) + bounded_weights, bounded_iterations, bounded_feasibility = ( + calibrate_entropy_bounded( + X, + y_target, + baseline_weights, + ss_values=ss_values, + ss_target=ss_target, + payroll_values=payroll_values, + payroll_target=payroll_target, + h6_income_values=h6_income_values, + h6_revenue_target=h6_revenue_target, + oasdi_tob_values=oasdi_tob_values, + oasdi_tob_target=oasdi_tob_target, + hi_tob_values=hi_tob_values, + hi_tob_target=hi_tob_target, + n_ages=n_ages, + max_constraint_error_pct=approximate_max_error_pct, + max_iters=max_iters * 10, + tol=max(tol, 1e-10), + warm_weights=warm_weights, + ) + ) + audit["approximate_solution_used"] = True + audit["approximation_method"] = "bounded_entropy" + audit["approximate_solution_error_pct"] = float( + bounded_feasibility["best_case_max_pct_error"] + ) + return bounded_weights, bounded_iterations, audit + except RuntimeError as bounded_error: + audit["approximate_entropy_error"] = str(bounded_error) + + if dense_lp_weights is not None and dense_lp_info is not None: + audit["lp_fallback_used"] = True + audit["approximate_solution_used"] = True + audit["approximation_method"] = "lp_blend" + audit["approximate_solution_error_pct"] = float( + dense_lp_info["best_case_max_pct_error"] + ) + audit["lp_blend_lambda"] = float(dense_lp_info["blend_lambda"]) + return dense_lp_weights, iterations, audit + audit["lp_fallback_used"] = True audit["approximate_solution_used"] = True audit["approximation_method"] = "lp_minimax" diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index 6174af33b..2475cc9ec 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -9,6 +9,7 @@ assess_nonnegative_feasibility, build_calibration_audit, calibrate_entropy, + calibrate_entropy_bounded, calibrate_weights, ) from policyengine_us_data.datasets.cps.long_term.calibration_artifacts import ( @@ -608,7 +609,29 @@ def test_entropy_calibration_produces_nonnegative_weights_and_hits_targets(): np.testing.assert_allclose(np.dot(payroll_values, weights), payroll_target, rtol=1e-8, atol=1e-8) -def test_entropy_calibration_can_fall_back_to_lp_approximate_solution(): +def test_bounded_entropy_calibration_returns_positive_approximate_weights(): + X = np.array( + [ + [1.0, 0.0, 1.0], + [0.0, 1.0, 1.0], + ] + ) + y_target = np.array([1.0, 1.0, 3.0]) + baseline_weights = np.array([1.0, 1.0]) + + weights, _, info = calibrate_entropy_bounded( + X=X, + y_target=y_target, + baseline_weights=baseline_weights, + n_ages=3, + max_constraint_error_pct=40.0, + ) + + assert info["best_case_max_pct_error"] <= 40.0 + assert np.all(weights > 0) + + +def test_entropy_calibration_prefers_bounded_entropy_over_lp_approximate_solution(): X = np.array( [ [1.0, 0.0, 1.0], @@ -629,9 +652,9 @@ def test_entropy_calibration_can_fall_back_to_lp_approximate_solution(): ) assert audit["approximate_solution_used"] is True - assert audit["approximation_method"] == "lp_minimax" + assert audit["approximation_method"] == "bounded_entropy" assert audit["approximate_solution_error_pct"] > 10.0 - assert np.all(weights >= 0) + assert np.all(weights > 0) def test_entropy_calibration_uses_lp_exact_fallback_even_before_approximate_window( @@ -644,7 +667,7 @@ def test_entropy_calibration_uses_lp_exact_fallback_even_before_approximate_wind ) monkeypatch.setattr( calibration_module, - "calibrate_entropy_approximate", + "calibrate_lp_minimax", lambda *args, **kwargs: ( np.array([1.0, 2.0]), 1, From e00f6e593afc7bb1f455d11a71cbf4690007158b Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 31 Mar 2026 15:57:50 -0400 Subject: [PATCH 15/31] Fix review issues in long-run calibration harness --- .../datasets/cps/long_term/calibration.py | 16 ++++++++++--- .../cps/long_term/calibration_artifacts.py | 5 ++++ .../cps/long_term/calibration_profiles.py | 5 +++- .../cps/long_term/projection_utils.py | 16 +++++++++++-- .../cps/long_term/run_household_projection.py | 23 ++++++++++++++++--- .../test_long_term_calibration_contract.py | 15 ++++++++++-- 6 files changed, 69 insertions(+), 11 deletions(-) diff --git a/policyengine_us_data/datasets/cps/long_term/calibration.py b/policyengine_us_data/datasets/cps/long_term/calibration.py index 7cb030f8c..fe6f1bbfe 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration.py @@ -503,10 +503,14 @@ def objective_and_gradient(z): best_max_error_pct = float("inf") for start in starts: + def objective_with_gradient(z): + objective, gradient, _ = objective_and_gradient(z) + return objective, gradient + result = optimize.minimize( - lambda z: objective_and_gradient(z)[0], + objective_with_gradient, start, - jac=lambda z: objective_and_gradient(z)[1], + jac=True, method="L-BFGS-B", bounds=[(0.0, None)] * len(start), options={"maxiter": max_iters, "ftol": tol}, @@ -984,6 +988,7 @@ def build_calibration_audit( negative_values = np.abs(weights[neg_mask]) positive_mask = weights > 0 weight_sum = float(np.sum(weights)) + abs_weight_sum = float(np.sum(np.abs(weights))) if weight_sum > 0: sorted_weights = np.sort(weights) top_10_weight_share_pct = float(sorted_weights[-10:].sum() / weight_sum * 100) @@ -1002,7 +1007,12 @@ def build_calibration_audit( { "age_max_pct_error": float(age_errors.max()), "negative_weight_count": int(neg_mask.sum()), - "negative_weight_pct": float(100 * neg_mask.sum() / len(weights)), + "negative_weight_household_pct": float(100 * neg_mask.sum() / len(weights)), + "negative_weight_pct": float( + 100 * negative_values.sum() / abs_weight_sum + ) + if abs_weight_sum > 0 + else 0.0, "largest_negative_weight": float(negative_values.max()) if negative_values.size else 0.0, "positive_weight_count": int(positive_mask.sum()), "positive_weight_pct": float(100 * positive_mask.sum() / len(weights)), diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py index d1b71c3fc..2f92fba52 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py @@ -158,6 +158,11 @@ def update_dataset_manifest( "age_max_pct_error": calibration_audit.get("age_max_pct_error"), "max_constraint_pct_error": calibration_audit.get("max_constraint_pct_error"), "negative_weight_pct": calibration_audit.get("negative_weight_pct"), + "negative_weight_household_pct": calibration_audit.get( + "negative_weight_household_pct" + ), + "validation_passed": calibration_audit.get("validation_passed"), + "validation_issue_count": len(calibration_audit.get("validation_issues", [])), } year_set = {int(value) for value in manifest.get("years", [])} diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py index 59f2c27d6..a0a49298b 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py @@ -217,7 +217,7 @@ def approximate_window_for_year( return max( profile.approximate_windows, key=lambda window: ( - float("-inf") if window.end_year is None else window.end_year, + float("inf") if window.end_year is None else window.end_year, window.max_constraint_error_pct, window.max_age_error_pct, ), @@ -237,6 +237,9 @@ def build_profile_from_flags( use_h6_reform: bool, use_tob: bool, ) -> CalibrationProfile: + if use_tob and not use_greg: + use_greg = True + if not use_greg: for profile in NAMED_PROFILES.values(): if ( diff --git a/policyengine_us_data/datasets/cps/long_term/projection_utils.py b/policyengine_us_data/datasets/cps/long_term/projection_utils.py index eb9e9aa44..390dcc690 100644 --- a/policyengine_us_data/datasets/cps/long_term/projection_utils.py +++ b/policyengine_us_data/datasets/cps/long_term/projection_utils.py @@ -1,5 +1,6 @@ import os import gc +import sys import numpy as np import h5py @@ -178,9 +179,20 @@ def create_household_year_h5(year, household_weights, base_dataset_path, output_ df[col_name_new] = uprated_values df.drop(columns=[col], inplace=True) else: + print( + f"Warning: uprating {var_name} for {year} returned " + f"{len(uprated_values)} rows instead of {len(df)}; " + "renaming the base-year column without recalculation.", + file=sys.stderr, + ) df.rename(columns={col: col_name_new}, inplace=True) - except: + except Exception as error: + print( + f"Warning: failed to uprate {var_name} for {year}: {error}; " + "renaming the base-year column without recalculation.", + file=sys.stderr, + ) df.rename(columns={col: col_name_new}, inplace=True) dataset = Dataset.from_dataframe(df, year) @@ -203,7 +215,7 @@ def create_household_year_h5(year, household_weights, base_dataset_path, output_ if values.dtype == np.object_: try: values = values.astype("S") - except: + except (TypeError, ValueError): continue data[variable][period] = values diff --git a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py index c1eeb60a7..4b7cf8b6d 100644 --- a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py +++ b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py @@ -3,7 +3,7 @@ Usage: - python run_household_projection.py [START_YEAR] [END_YEAR] [--profile PROFILE] [--target-source SOURCE] [--output-dir DIR] [--save-h5] + python run_household_projection.py [START_YEAR] [END_YEAR] [--profile PROFILE] [--target-source SOURCE] [--output-dir DIR] [--save-h5] [--allow-validation-failures] python run_household_projection.py [START_YEAR] [END_YEAR] [--greg] [--use-ss] [--use-payroll] [--use-h6-reform] [--use-tob] [--save-h5] START_YEAR: Optional starting year (default: 2025) @@ -11,6 +11,7 @@ --profile: Named calibration contract (recommended) --target-source: Named long-term target source package --output-dir: Output directory for generated H5 files and metadata + --allow-validation-failures: Record validation issues in metadata and continue instead of aborting the run --greg: Use GREG calibration instead of IPF (optional) --use-ss: Include Social Security benefit totals as calibration target (requires --greg) --use-payroll: Include taxable payroll totals as calibration target (requires --greg) @@ -269,6 +270,13 @@ def create_h6_reform(): OUTPUT_DIR = sys.argv[output_dir_index + 1] del sys.argv[output_dir_index : output_dir_index + 2] +ALLOW_VALIDATION_FAILURES = "--allow-validation-failures" in sys.argv +if ALLOW_VALIDATION_FAILURES: + sys.argv.remove("--allow-validation-failures") +ALLOW_VALIDATION_FAILURES = ALLOW_VALIDATION_FAILURES or ( + os.environ.get("PEUD_ALLOW_INVALID_ARTIFACTS", "").lower() in {"1", "true", "yes"} +) + USE_GREG = "--greg" in sys.argv if USE_GREG: @@ -673,17 +681,26 @@ def create_h6_reform(): PROFILE, year=year, ) + calibration_audit["validation_issues"] = validation_issues + calibration_audit["validation_passed"] = not bool(validation_issues) if validation_issues: issue_text = "; ".join(validation_issues) - raise RuntimeError(f"Calibration validation failed for {year}: {issue_text}") + if not ALLOW_VALIDATION_FAILURES: + raise RuntimeError(f"Calibration validation failed for {year}: {issue_text}") + print( + f" [WARN {year}] Validation issues recorded but not fatal: {issue_text}", + file=sys.stderr, + ) if year in display_years and CALIBRATION_METHOD in {"greg", "entropy"}: n_neg = calibration_audit["negative_weight_count"] if n_neg > 0: pct_neg = calibration_audit["negative_weight_pct"] + hh_pct_neg = calibration_audit.get("negative_weight_household_pct", 0.0) max_neg = calibration_audit["largest_negative_weight"] print( - f" [DEBUG {year}] Negative weights: {n_neg} ({pct_neg:.2f}%), " + f" [DEBUG {year}] Negative weights: {n_neg} households " + f"({hh_pct_neg:.2f}% of households, {pct_neg:.2f}% of weight mass), " f"largest: {max_neg:,.0f}" ) else: diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index 2475cc9ec..c3385956c 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -19,6 +19,7 @@ write_year_metadata, ) from policyengine_us_data.datasets.cps.long_term.calibration_profiles import ( + approximate_window_for_year, build_profile_from_flags, classify_calibration_quality, get_profile, @@ -277,8 +278,17 @@ def test_legacy_flags_map_to_named_profile(): use_h6_reform=False, use_tob=True, ) - assert profile.name == "custom-ss-payroll-tob" - assert profile.calibration_method == "ipf" + assert profile.name == "custom-greg-ss-payroll-tob" + assert profile.calibration_method == "greg" + assert profile.use_greg is True + + +def test_approximate_window_none_selects_open_ended_tail(): + profile = get_profile("ss-payroll-tob") + window = approximate_window_for_year(profile, None) + assert window is not None + assert window.start_year == 2096 + assert window.end_year is None def test_strict_greg_failure_raises(): @@ -322,6 +332,7 @@ def test_build_calibration_audit_reports_constraint_error(): assert audit["constraints"]["payroll_total"]["pct_error"] == -50.0 assert audit["positive_weight_count"] == 2 assert audit["positive_weight_pct"] == 100.0 + assert audit["negative_weight_household_pct"] == 0.0 assert audit["effective_sample_size"] == pytest.approx(2.0) assert audit["top_10_weight_share_pct"] == pytest.approx(100.0) assert audit["top_100_weight_share_pct"] == pytest.approx(100.0) From 903618378b6f48c9c3daf28d4d833165a01c1b91 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 31 Mar 2026 17:26:44 -0400 Subject: [PATCH 16/31] Refine late-tail calibration metadata and caching --- .../datasets/cps/long_term/calibration.py | 154 ++++++++----- .../cps/long_term/calibration_artifacts.py | 53 ++++- .../test_long_term_calibration_contract.py | 210 ++++++++++++++++-- 3 files changed, 332 insertions(+), 85 deletions(-) diff --git a/policyengine_us_data/datasets/cps/long_term/calibration.py b/policyengine_us_data/datasets/cps/long_term/calibration.py index fe6f1bbfe..238350b4d 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration.py @@ -279,7 +279,9 @@ def calibrate_entropy( baseline_weights = np.asarray(baseline_weights, dtype=float) gram = A_scaled.T @ (baseline_weights[:, None] * A_scaled) gram += np.eye(gram.shape[0]) * 1e-12 - beta0 = np.linalg.solve(gram, targets_scaled - (A_scaled.T @ baseline_weights)) + beta0 = np.linalg.solve( + gram, targets_scaled - (A_scaled.T @ baseline_weights) + ) def objective_gradient_hessian(beta): eta = np.clip(A_scaled @ beta, -700, 700) @@ -291,10 +293,19 @@ def objective_gradient_hessian(beta): return objective, gradient, hessian def solve_with_root(beta_start): + _cache = {} + + def _cached_ogh(z): + key = z.tobytes() + if key not in _cache: + _cache.clear() + _cache[key] = objective_gradient_hessian(z) + return _cache[key] + result = optimize.root( - lambda z: objective_gradient_hessian(z)[1], + lambda z: _cached_ogh(z)[1], beta_start, - jac=lambda z: objective_gradient_hessian(z)[2], + jac=lambda z: _cached_ogh(z)[2], method="hybr", options={"xtol": tol}, ) @@ -302,7 +313,11 @@ def solve_with_root(beta_start): return None _, gradient, _ = objective_gradient_hessian(result.x) max_error = float( - np.max(100 * np.abs(gradient) / np.maximum(np.abs(targets_scaled), 1e-12)) + np.max( + 100 + * np.abs(gradient) + / np.maximum(np.abs(targets_scaled), 1e-12) + ) ) if max_error > tol * 100: return None @@ -328,7 +343,11 @@ def infeasibility_error(prefix): for iterations in range(1, max_iters + 1): objective, gradient, hessian = objective_gradient_hessian(beta) final_max_error = float( - np.max(100 * np.abs(gradient) / np.maximum(np.abs(targets_scaled), 1e-12)) + np.max( + 100 + * np.abs(gradient) + / np.maximum(np.abs(targets_scaled), 1e-12) + ) ) if final_max_error <= tol * 100: break @@ -342,7 +361,9 @@ def infeasibility_error(prefix): step = 1.0 while step >= 1e-8: candidate = beta - step * delta - candidate_objective, candidate_gradient, _ = objective_gradient_hessian(candidate) + candidate_objective, candidate_gradient, _ = ( + objective_gradient_hessian(candidate) + ) candidate_max_error = float( np.max( 100 @@ -456,7 +477,9 @@ def objective_and_gradient(z): exp_eta = np.exp(eta) weights = baseline_weights * exp_eta achieved = A_scaled.T @ weights - objective = float(np.sum(weights) + upper_bounds @ alpha - lower_bounds @ gamma) + objective = float( + np.sum(weights) + upper_bounds @ alpha - lower_bounds @ gamma + ) gradient = np.concatenate( [ upper_bounds - achieved, @@ -498,15 +521,15 @@ def objective_and_gradient(z): ) ) + def objective_with_gradient(z): + objective, gradient, _ = objective_and_gradient(z) + return objective, gradient + best_result = None best_weights = None best_max_error_pct = float("inf") for start in starts: - def objective_with_gradient(z): - objective, gradient, _ = objective_and_gradient(z) - return objective, gradient - result = optimize.minimize( objective_with_gradient, start, @@ -518,9 +541,7 @@ def objective_with_gradient(z): objective, gradient, weights = objective_and_gradient(result.x) achieved = A_scaled.T @ weights - max_error_pct = float( - np.max(np.abs(achieved - targets_scaled)) * 100 - ) + max_error_pct = float(np.max(np.abs(achieved - targets_scaled)) * 100) if max_error_pct < best_max_error_pct: best_result = result @@ -540,7 +561,9 @@ def objective_with_gradient(z): ) if best_result is None or best_weights is None: - raise RuntimeError("Approximate bounded entropy calibration did not run.") + raise RuntimeError( + "Approximate bounded entropy calibration did not run." + ) raise RuntimeError( "Approximate bounded entropy calibration failed: " @@ -587,7 +610,8 @@ def densify_lp_solution( lam = (lo + hi) / 2.0 candidate_weights = (1.0 - lam) * lp_weights + lam * baseline_weights candidate_error_pct = float( - np.max(np.abs(A_scaled.T @ candidate_weights - targets_scaled)) * 100 + np.max(np.abs(A_scaled.T @ candidate_weights - targets_scaled)) + * 100 ) if candidate_error_pct <= max_constraint_error_pct + 1e-6: best_lambda = lam @@ -600,6 +624,7 @@ def densify_lp_solution( return best_weights, { "blend_lambda": best_lambda, "best_case_max_pct_error": best_error_pct, + "densification_effective": best_lambda > 0.0, } @@ -645,7 +670,9 @@ def calibrate_lp_minimax( A = aux_df.to_numpy(dtype=float) targets = np.array(list(controls.values()), dtype=float) - feasibility = assess_nonnegative_feasibility(A, targets, return_weights=True) + feasibility = assess_nonnegative_feasibility( + A, targets, return_weights=True + ) weights = feasibility.get("weights") if not feasibility["success"] or weights is None: raise RuntimeError( @@ -679,7 +706,9 @@ def assess_nonnegative_feasibility(A, targets, *, return_weights=False): b_rel = targets / scales constraint_matrix = sparse.csr_matrix(A_rel) - epsilon_column = sparse.csc_matrix(np.ones((constraint_matrix.shape[0], 1))) + epsilon_column = sparse.csc_matrix( + np.ones((constraint_matrix.shape[0], 1)) + ) A_ub = sparse.vstack( [ sparse.hstack([constraint_matrix, -epsilon_column]), @@ -857,7 +886,9 @@ def calibrate_weights( hi_tob_target=hi_tob_target, n_ages=n_ages, ) - approximate_error_pct = float(feasibility["best_case_max_pct_error"]) + approximate_error_pct = float( + feasibility["best_case_max_pct_error"] + ) if approximate_error_pct <= max(tol * 100, 1e-6): audit["lp_fallback_used"] = True audit["approximation_method"] = "lp_minimax_exact" @@ -910,27 +941,29 @@ def calibrate_weights( warm_weights = [w_new] if dense_lp_weights is not None: warm_weights.insert(0, dense_lp_weights) - bounded_weights, bounded_iterations, bounded_feasibility = ( - calibrate_entropy_bounded( - X, - y_target, - baseline_weights, - ss_values=ss_values, - ss_target=ss_target, - payroll_values=payroll_values, - payroll_target=payroll_target, - h6_income_values=h6_income_values, - h6_revenue_target=h6_revenue_target, - oasdi_tob_values=oasdi_tob_values, - oasdi_tob_target=oasdi_tob_target, - hi_tob_values=hi_tob_values, - hi_tob_target=hi_tob_target, - n_ages=n_ages, - max_constraint_error_pct=approximate_max_error_pct, - max_iters=max_iters * 10, - tol=max(tol, 1e-10), - warm_weights=warm_weights, - ) + ( + bounded_weights, + bounded_iterations, + bounded_feasibility, + ) = calibrate_entropy_bounded( + X, + y_target, + baseline_weights, + ss_values=ss_values, + ss_target=ss_target, + payroll_values=payroll_values, + payroll_target=payroll_target, + h6_income_values=h6_income_values, + h6_revenue_target=h6_revenue_target, + oasdi_tob_values=oasdi_tob_values, + oasdi_tob_target=oasdi_tob_target, + hi_tob_values=hi_tob_values, + hi_tob_target=hi_tob_target, + n_ages=n_ages, + max_constraint_error_pct=approximate_max_error_pct, + max_iters=max_iters * 10, + tol=max(tol, 1e-10), + warm_weights=warm_weights, ) audit["approximate_solution_used"] = True audit["approximation_method"] = "bounded_entropy" @@ -944,7 +977,10 @@ def calibrate_weights( if dense_lp_weights is not None and dense_lp_info is not None: audit["lp_fallback_used"] = True audit["approximate_solution_used"] = True - audit["approximation_method"] = "lp_blend" + densified = dense_lp_info.get("densification_effective", False) + audit["approximation_method"] = ( + "lp_blend" if densified else "lp_minimax" + ) audit["approximate_solution_error_pct"] = float( dense_lp_info["best_case_max_pct_error"] ) @@ -982,7 +1018,11 @@ def build_calibration_audit( hi_tob_target=None, ): achieved_ages = X.T @ weights - age_errors = np.abs(achieved_ages - y_target) / np.maximum(np.abs(y_target), 1e-10) * 100 + age_errors = ( + np.abs(achieved_ages - y_target) + / np.maximum(np.abs(y_target), 1e-10) + * 100 + ) neg_mask = weights < 0 negative_values = np.abs(weights[neg_mask]) @@ -991,8 +1031,12 @@ def build_calibration_audit( abs_weight_sum = float(np.sum(np.abs(weights))) if weight_sum > 0: sorted_weights = np.sort(weights) - top_10_weight_share_pct = float(sorted_weights[-10:].sum() / weight_sum * 100) - top_100_weight_share_pct = float(sorted_weights[-100:].sum() / weight_sum * 100) + top_10_weight_share_pct = float( + sorted_weights[-10:].sum() / weight_sum * 100 + ) + top_100_weight_share_pct = float( + sorted_weights[-100:].sum() / weight_sum * 100 + ) else: top_10_weight_share_pct = 0.0 top_100_weight_share_pct = 0.0 @@ -1007,15 +1051,21 @@ def build_calibration_audit( { "age_max_pct_error": float(age_errors.max()), "negative_weight_count": int(neg_mask.sum()), - "negative_weight_household_pct": float(100 * neg_mask.sum() / len(weights)), - "negative_weight_pct": float( - 100 * negative_values.sum() / abs_weight_sum - ) - if abs_weight_sum > 0 - else 0.0, - "largest_negative_weight": float(negative_values.max()) if negative_values.size else 0.0, + "negative_weight_household_pct": float( + 100 * neg_mask.sum() / len(weights) + ), + "negative_weight_pct": ( + float(100 * negative_values.sum() / abs_weight_sum) + if abs_weight_sum > 0 + else 0.0 + ), + "largest_negative_weight": ( + float(negative_values.max()) if negative_values.size else 0.0 + ), "positive_weight_count": int(positive_mask.sum()), - "positive_weight_pct": float(100 * positive_mask.sum() / len(weights)), + "positive_weight_pct": float( + 100 * positive_mask.sum() / len(weights) + ), "effective_sample_size": effective_sample_size, "top_10_weight_share_pct": top_10_weight_share_pct, "top_100_weight_share_pct": top_100_weight_share_pct, diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py index 2f92fba52..5be53d18d 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py @@ -6,9 +6,17 @@ from typing import Any try: - from .calibration_profiles import classify_calibration_quality, get_profile + from .calibration_profiles import ( + classify_calibration_quality, + get_profile, + validate_calibration_audit, + ) except ImportError: # pragma: no cover - script execution fallback - from calibration_profiles import classify_calibration_quality, get_profile + from calibration_profiles import ( + classify_calibration_quality, + get_profile, + validate_calibration_audit, + ) CONTRACT_VERSION = 1 @@ -29,7 +37,13 @@ def normalize_metadata(metadata: dict[str, Any]) -> dict[str, Any]: if "max_constraint_pct_error" not in audit: audit["max_constraint_pct_error"] = float( - max((abs(stats.get("pct_error", 0.0)) for stats in constraints.values()), default=0.0) + max( + ( + abs(stats.get("pct_error", 0.0)) + for stats in constraints.values() + ), + default=0.0, + ) ) if audit.get("lp_fallback_used"): @@ -63,6 +77,20 @@ def normalize_metadata(metadata: dict[str, Any]) -> dict[str, Any]: audit["approximation_method"] = "lp_minimax" audit["approximate_solution_used"] = True + if "validation_passed" not in audit and profile_data.get("name"): + try: + profile = get_profile(profile_data["name"]) + except ValueError: + profile = None + if profile is not None: + issues = validate_calibration_audit( + audit, + profile, + year=normalized.get("year"), + ) + audit["validation_passed"] = not bool(issues) + audit.setdefault("validation_issues", issues) + return normalized @@ -129,10 +157,10 @@ def update_dataset_manifest( ) manifest_profile = json.loads(json.dumps(manifest["profile"])) if manifest_profile != profile: - if ( - manifest_profile.get("name") == profile.get("name") - and manifest_profile.get("calibration_method") - == profile.get("calibration_method") + if manifest_profile.get("name") == profile.get( + "name" + ) and manifest_profile.get("calibration_method") == profile.get( + "calibration_method" ): manifest["profile"] = profile else: @@ -156,13 +184,17 @@ def update_dataset_manifest( "method_used": calibration_audit.get("method_used"), "fell_back_to_ipf": calibration_audit.get("fell_back_to_ipf"), "age_max_pct_error": calibration_audit.get("age_max_pct_error"), - "max_constraint_pct_error": calibration_audit.get("max_constraint_pct_error"), + "max_constraint_pct_error": calibration_audit.get( + "max_constraint_pct_error" + ), "negative_weight_pct": calibration_audit.get("negative_weight_pct"), "negative_weight_household_pct": calibration_audit.get( "negative_weight_household_pct" ), "validation_passed": calibration_audit.get("validation_passed"), - "validation_issue_count": len(calibration_audit.get("validation_issues", [])), + "validation_issue_count": len( + calibration_audit.get("validation_issues", []) + ), } year_set = {int(value) for value in manifest.get("years", [])} @@ -173,6 +205,9 @@ def update_dataset_manifest( "end": max(year_set), } manifest["generated_at"] = datetime.now(timezone.utc).isoformat() + manifest["contains_invalid_artifacts"] = any( + entry.get("validation_passed") is False for entry in datasets.values() + ) manifest_path.write_text( json.dumps(manifest, indent=2, sort_keys=True) + "\n", diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index c3385956c..e6e31de83 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -4,7 +4,9 @@ import numpy as np import pytest -from policyengine_us_data.datasets.cps.long_term import calibration as calibration_module +from policyengine_us_data.datasets.cps.long_term import ( + calibration as calibration_module, +) from policyengine_us_data.datasets.cps.long_term.calibration import ( assess_nonnegative_feasibility, build_calibration_audit, @@ -69,12 +71,30 @@ def _toy_support_dataframe(): "age__2024": [70.0, 68.0, 80.0, 77.0, 60.0], "household_weight__2024": [10.0, 10.0, 8.0, 8.0, 5.0], "person_weight__2024": [10.0, 10.0, 8.0, 8.0, 5.0], - "social_security_retirement__2024": [20_000.0, 0.0, 30_000.0, 0.0, 0.0], + "social_security_retirement__2024": [ + 20_000.0, + 0.0, + 30_000.0, + 0.0, + 0.0, + ], "social_security_disability__2024": [0.0, 0.0, 0.0, 0.0, 0.0], "social_security_survivors__2024": [0.0, 0.0, 0.0, 0.0, 0.0], "social_security_dependents__2024": [0.0, 0.0, 0.0, 0.0, 0.0], - "employment_income_before_lsr__2024": [5_000.0, 0.0, 12_000.0, 0.0, 50_000.0], - "self_employment_income_before_lsr__2024": [0.0, 0.0, 0.0, 0.0, 0.0], + "employment_income_before_lsr__2024": [ + 5_000.0, + 0.0, + 12_000.0, + 0.0, + 50_000.0, + ], + "self_employment_income_before_lsr__2024": [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ], } ) ) @@ -140,9 +160,15 @@ def test_support_augmentation_clones_households_with_new_ids(): ) assert report["base_household_count"] == 3 assert report["augmented_household_count"] == 4 - cloned_household_ids = set(augmented_df["household_id__2024"].unique()) - {1, 2, 3} + cloned_household_ids = set(augmented_df["household_id__2024"].unique()) - { + 1, + 2, + 3, + } assert len(cloned_household_ids) == 1 - cloned_rows = augmented_df[augmented_df["household_id__2024"].isin(cloned_household_ids)] + cloned_rows = augmented_df[ + augmented_df["household_id__2024"].isin(cloned_household_ids) + ] assert cloned_rows["age__2024"].max() == pytest.approx(80.0) assert cloned_rows["household_weight__2024"].iloc[0] == pytest.approx(5.0) assert cloned_rows["person_id__2024"].min() > df["person_id__2024"].max() @@ -178,12 +204,22 @@ def test_support_augmentation_synthesizes_composite_payroll_household(): ) assert report["base_household_count"] == 3 assert report["augmented_household_count"] == 4 - cloned_household_ids = set(augmented_df["household_id__2024"].unique()) - {1, 2, 3} + cloned_household_ids = set(augmented_df["household_id__2024"].unique()) - { + 1, + 2, + 3, + } assert len(cloned_household_ids) == 1 - cloned_rows = augmented_df[augmented_df["household_id__2024"].isin(cloned_household_ids)] + cloned_rows = augmented_df[ + augmented_df["household_id__2024"].isin(cloned_household_ids) + ] assert cloned_rows["age__2024"].max() == pytest.approx(80.0) - assert cloned_rows["social_security_retirement__2024"].sum() == pytest.approx(30_000.0) - assert cloned_rows["employment_income_before_lsr__2024"].sum() == pytest.approx(37_000.0) + assert cloned_rows[ + "social_security_retirement__2024" + ].sum() == pytest.approx(30_000.0) + assert cloned_rows[ + "employment_income_before_lsr__2024" + ].sum() == pytest.approx(37_000.0) def test_support_augmentation_appends_single_person_synthetic_grid_households(): @@ -218,13 +254,13 @@ def test_support_augmentation_appends_single_person_synthetic_grid_households(): name="grid-profile", description="Toy single-person synthetic grid.", rules=( - SinglePersonSyntheticGridRule( - name="older_grid", - template_min_max_age=75, - template_max_max_age=86, - target_ages=(77, 85), - ss_quantiles=(0.5,), - payroll_quantiles=(0.5,), + SinglePersonSyntheticGridRule( + name="older_grid", + template_min_max_age=75, + template_max_max_age=86, + target_ages=(77, 85), + ss_quantiles=(0.5,), + payroll_quantiles=(0.5,), template_ss_state="positive", template_payroll_state="any", payroll_donor_min_max_age=55, @@ -240,12 +276,20 @@ def test_support_augmentation_appends_single_person_synthetic_grid_households(): ) assert report["base_household_count"] == 3 assert report["augmented_household_count"] == 5 - synthetic_household_ids = set(augmented_df["household_id__2024"].unique()) - {1, 2, 3} + synthetic_household_ids = set( + augmented_df["household_id__2024"].unique() + ) - {1, 2, 3} assert len(synthetic_household_ids) == 2 - synthetic_rows = augmented_df[augmented_df["household_id__2024"].isin(synthetic_household_ids)] + synthetic_rows = augmented_df[ + augmented_df["household_id__2024"].isin(synthetic_household_ids) + ] assert set(synthetic_rows["age__2024"].tolist()) == {77.0, 85.0} - assert set(synthetic_rows["social_security_retirement__2024"].tolist()) == {22_000.0} - assert set(synthetic_rows["employment_income_before_lsr__2024"].tolist()) == {50_000.0} + assert set( + synthetic_rows["social_security_retirement__2024"].tolist() + ) == {22_000.0} + assert set( + synthetic_rows["employment_income_before_lsr__2024"].tolist() + ) == {50_000.0} def test_age_bin_helpers_preserve_population_totals(): @@ -617,7 +661,9 @@ def test_entropy_calibration_produces_nonnegative_weights_and_hits_targets(): assert np.all(weights > 0) np.testing.assert_allclose(X.T @ weights, y_target, rtol=1e-8, atol=1e-8) - np.testing.assert_allclose(np.dot(payroll_values, weights), payroll_target, rtol=1e-8, atol=1e-8) + np.testing.assert_allclose( + np.dot(payroll_values, weights), payroll_target, rtol=1e-8, atol=1e-8 + ) def test_bounded_entropy_calibration_returns_positive_approximate_weights(): @@ -674,7 +720,9 @@ def test_entropy_calibration_uses_lp_exact_fallback_even_before_approximate_wind monkeypatch.setattr( calibration_module, "calibrate_entropy", - lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("entropy stalled")), + lambda *args, **kwargs: (_ for _ in ()).throw( + RuntimeError("entropy stalled") + ), ) monkeypatch.setattr( calibration_module, @@ -722,7 +770,9 @@ def test_nonnegative_feasibility_diagnostic_distinguishes_feasible_and_infeasibl ] ) infeasible_targets = np.array([1.0, 1.0, 3.0]) - infeasible = assess_nonnegative_feasibility(infeasible_A, infeasible_targets) + infeasible = assess_nonnegative_feasibility( + infeasible_A, infeasible_targets + ) assert infeasible["success"] is True assert infeasible["best_case_max_pct_error"] > 10.0 @@ -739,3 +789,115 @@ def test_long_term_target_sources_are_available_and_distinct(): source_name="trustees_2025_current_law", ) assert payroll_2026 == pytest.approx(11_129_000_000_000.0) + + +def test_normalize_metadata_backfills_validation_passed(): + metadata = normalize_metadata( + { + "year": 2091, + "profile": {"name": "ss-payroll-tob"}, + "calibration_audit": { + "lp_fallback_used": True, + "approximation_method": "lp_blend", + "approximate_solution_error_pct": 16.0, + "max_constraint_pct_error": 16.0, + "age_max_pct_error": 14.5, + "negative_weight_pct": 0.0, + "positive_weight_count": 6840, + "effective_sample_size": 12.0, + "top_10_weight_share_pct": 80.0, + "top_100_weight_share_pct": 99.0, + "constraints": { + "ss_total": {"pct_error": 14.5}, + "payroll_total": {"pct_error": 16.0}, + }, + }, + } + ) + + audit = metadata["calibration_audit"] + assert audit["validation_passed"] is False + assert isinstance(audit["validation_issues"], list) + assert len(audit["validation_issues"]) > 0 + + +def test_manifest_contains_invalid_artifacts_flag(tmp_path): + profile = get_profile("ss-payroll-tob") + + valid_audit = { + "method_used": "entropy", + "fell_back_to_ipf": False, + "age_max_pct_error": 0.0, + "negative_weight_pct": 0.0, + "positive_weight_count": 70000, + "effective_sample_size": 5000.0, + "top_10_weight_share_pct": 1.5, + "top_100_weight_share_pct": 10.0, + "max_constraint_pct_error": 0.0, + "constraints": {}, + "validation_passed": True, + "validation_issues": [], + } + + invalid_audit = { + "method_used": "entropy", + "fell_back_to_ipf": False, + "age_max_pct_error": 14.0, + "negative_weight_pct": 0.0, + "positive_weight_count": 6840, + "effective_sample_size": 12.0, + "top_10_weight_share_pct": 80.0, + "top_100_weight_share_pct": 99.0, + "max_constraint_pct_error": 16.0, + "constraints": {"payroll_total": {"pct_error": 16.0}}, + "validation_passed": False, + "validation_issues": ["ESS too low"], + } + + # First year: valid + year_2030 = tmp_path / "2030.h5" + year_2030.write_text("", encoding="utf-8") + metadata_2030 = write_year_metadata( + year_2030, + year=2030, + base_dataset_path="test.h5", + profile=profile.to_dict(), + calibration_audit=valid_audit, + ) + manifest_path = update_dataset_manifest( + tmp_path, + year=2030, + h5_path=year_2030, + metadata_path=metadata_2030, + base_dataset_path="test.h5", + profile=profile.to_dict(), + calibration_audit=valid_audit, + ) + + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + assert manifest["contains_invalid_artifacts"] is False + + # Second year: invalid + year_2091 = tmp_path / "2091.h5" + year_2091.write_text("", encoding="utf-8") + metadata_2091 = write_year_metadata( + year_2091, + year=2091, + base_dataset_path="test.h5", + profile=profile.to_dict(), + calibration_audit=invalid_audit, + ) + update_dataset_manifest( + tmp_path, + year=2091, + h5_path=year_2091, + metadata_path=metadata_2091, + base_dataset_path="test.h5", + profile=profile.to_dict(), + calibration_audit=invalid_audit, + ) + + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + assert manifest["contains_invalid_artifacts"] is True + assert manifest["datasets"]["2030"]["validation_passed"] is True + assert manifest["datasets"]["2091"]["validation_passed"] is False From 79858d6c05dbceb5dafd288d513932e78d9bbbd0 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 31 Mar 2026 19:02:07 -0400 Subject: [PATCH 17/31] Add publishable horizon assessment tool --- .../long_term/assess_publishable_horizon.py | 443 ++++++++++++++++++ 1 file changed, 443 insertions(+) create mode 100644 policyengine_us_data/datasets/cps/long_term/assess_publishable_horizon.py diff --git a/policyengine_us_data/datasets/cps/long_term/assess_publishable_horizon.py b/policyengine_us_data/datasets/cps/long_term/assess_publishable_horizon.py new file mode 100644 index 000000000..3c482243c --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/assess_publishable_horizon.py @@ -0,0 +1,443 @@ +from __future__ import annotations + +import argparse +import csv +import gc +from pathlib import Path +import re +import sys + +import numpy as np + +from policyengine_us import Microsimulation + +from calibration import build_calibration_audit, calibrate_weights +from calibration_profiles import ( + approximate_window_for_year, + classify_calibration_quality, + get_profile, + validate_calibration_audit, +) +from projection_utils import ( + aggregate_age_targets, + aggregate_household_age_matrix, + build_age_bins, + build_household_age_matrix, +) +from ssa_data import ( + get_long_term_target_source, + load_hi_tob_projections, + load_oasdi_tob_projections, + load_ssa_age_projections, + load_ssa_benefit_projections, + load_taxable_payroll_projections, + set_long_term_target_source, +) + +try: + from samplics.weighting import SampleWeight +except ImportError: # pragma: no cover - only needed for greg profiles + SampleWeight = None + + +DEFAULT_BASE_DATASET_PATH = ( + "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5" +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Assess publishable microsimulation horizon quality for selected years." + ) + ) + parser.add_argument( + "--profile", + default="ss-payroll-tob", + help="Named calibration profile to assess.", + ) + parser.add_argument( + "--target-source", + default=get_long_term_target_source(), + help="Named long-run target source package.", + ) + parser.add_argument( + "--years", + default="2075,2080,2085,2090,2095,2100", + help="Comma-separated years to assess.", + ) + parser.add_argument( + "--base-dataset", + default=DEFAULT_BASE_DATASET_PATH, + help="Base microsimulation dataset path.", + ) + parser.add_argument( + "--output", + type=Path, + help="Optional CSV output path. Defaults to stdout.", + ) + return parser.parse_args() + + +def parse_years(raw: str) -> list[int]: + years = [int(value.strip()) for value in raw.split(",") if value.strip()] + if not years: + raise ValueError("At least one year must be provided.") + return sorted(set(years)) + + +def maybe_build_calibrator(method: str): + if method != "greg": + return None + if SampleWeight is None: + raise ImportError( + "samplics is required for GREG calibration. " + "Install with: pip install policyengine-us-data[calibration]" + ) + return SampleWeight() + + +def benchmark_tob_values( + year: int, + weights: np.ndarray, + *, + oasdi_tob_values: np.ndarray | None, + hi_tob_values: np.ndarray | None, +) -> dict[str, float] | None: + if oasdi_tob_values is None or hi_tob_values is None: + return None + + oasdi_target = float(load_oasdi_tob_projections(year)) + oasdi_achieved = float(np.sum(oasdi_tob_values * weights)) + hi_target = float(load_hi_tob_projections(year)) + hi_achieved = float(np.sum(hi_tob_values * weights)) + + return { + "oasdi_tob_benchmark_pct_error": ( + 0.0 + if oasdi_target == 0 + else (oasdi_achieved - oasdi_target) / oasdi_target * 100 + ), + "hi_tob_benchmark_pct_error": ( + 0.0 + if hi_target == 0 + else (hi_achieved - hi_target) / hi_target * 100 + ), + } + + +def assess_years( + *, + years: list[int], + profile_name: str, + target_source: str, + base_dataset_path: str, +) -> list[dict[str, object]]: + profile = get_profile(profile_name) + if profile.use_h6_reform: + raise NotImplementedError( + "assess_publishable_horizon.py does not yet support H6-calibrated profiles." + ) + + set_long_term_target_source(target_source) + calibrator = maybe_build_calibrator(profile.calibration_method) + + start_year = min(years) + end_year = max(years) + target_matrix = load_ssa_age_projections(start_year=start_year, end_year=end_year) + n_ages = target_matrix.shape[0] + + sim = Microsimulation(dataset=base_dataset_path) + X, _, _ = build_household_age_matrix(sim, n_ages) + del sim + gc.collect() + + aggregated_age_cache: dict[int, tuple[np.ndarray, np.ndarray]] = {} + rows: list[dict[str, object]] = [] + + for year in years: + print(f"[assess_publishable_horizon] year={year}", file=sys.stderr, flush=True) + year_idx = year - start_year + sim = Microsimulation(dataset=base_dataset_path) + + household_microseries = sim.calculate("household_id", map_to="household") + baseline_weights = household_microseries.weights.values + + ss_values = None + ss_target = None + if profile.use_ss: + ss_values = sim.calculate( + "social_security", + period=year, + map_to="household", + ).values + ss_target = load_ssa_benefit_projections(year) + + payroll_values = None + payroll_target = None + if profile.use_payroll: + taxable_wages = sim.calculate( + "taxable_earnings_for_social_security", + period=year, + map_to="household", + ).values + taxable_self_employment = sim.calculate( + "social_security_taxable_self_employment_income", + period=year, + map_to="household", + ).values + payroll_values = taxable_wages + taxable_self_employment + payroll_target = load_taxable_payroll_projections(year) + + oasdi_tob_values = None + hi_tob_values = None + if profile.use_tob or profile.benchmark_tob: + oasdi_tob_values = sim.calculate( + "tob_revenue_oasdi", + period=year, + map_to="household", + ).values + hi_tob_values = sim.calculate( + "tob_revenue_medicare_hi", + period=year, + map_to="household", + ).values + + approximate_window = approximate_window_for_year(profile, year) + age_bucket_size = ( + approximate_window.age_bucket_size + if approximate_window is not None and approximate_window.age_bucket_size + else 1 + ) + + if age_bucket_size > 1: + if age_bucket_size not in aggregated_age_cache: + age_bins = build_age_bins(n_ages=n_ages, bucket_size=age_bucket_size) + aggregated_age_cache[age_bucket_size] = ( + aggregate_household_age_matrix(X, age_bins), + aggregate_age_targets(target_matrix, age_bins), + ) + X_current, aggregated_target_matrix = aggregated_age_cache[age_bucket_size] + y_target = aggregated_target_matrix[:, year_idx] + else: + X_current = X + y_target = target_matrix[:, year_idx] + + try: + weights, iterations, calibration_event = calibrate_weights( + X=X_current, + y_target=y_target, + baseline_weights=baseline_weights, + method=profile.calibration_method, + calibrator=calibrator, + ss_values=ss_values, + ss_target=ss_target, + payroll_values=payroll_values, + payroll_target=payroll_target, + h6_income_values=None, + h6_revenue_target=None, + oasdi_tob_values=oasdi_tob_values if profile.use_tob else None, + oasdi_tob_target=load_oasdi_tob_projections(year) + if profile.use_tob + else None, + hi_tob_values=hi_tob_values if profile.use_tob else None, + hi_tob_target=load_hi_tob_projections(year) + if profile.use_tob + else None, + n_ages=X_current.shape[1], + max_iters=100, + tol=1e-6, + verbose=False, + allow_fallback_to_ipf=profile.allow_greg_fallback, + allow_approximate_entropy=approximate_window is not None, + approximate_max_error_pct=( + approximate_window.max_constraint_error_pct + if approximate_window is not None + else None + ), + ) + except RuntimeError as error: + row: dict[str, object] = { + "year": year, + "target_source": target_source, + "profile": profile.name, + "calibration_quality": "failed", + "approximation_method": "runtime_error", + "iterations": None, + "age_bucket_size": age_bucket_size, + "window_max_constraint_error_pct": ( + approximate_window.max_constraint_error_pct + if approximate_window is not None + else profile.max_constraint_error_pct + ), + "window_max_age_error_pct": ( + approximate_window.max_age_error_pct + if approximate_window is not None + else profile.max_age_error_pct + ), + "max_constraint_pct_error": None, + "age_max_pct_error": None, + "positive_weight_count": None, + "effective_sample_size": None, + "top_10_weight_share_pct": None, + "top_100_weight_share_pct": None, + "negative_weight_pct": None, + "validation_passed": False, + "validation_issue_count": 1, + "validation_issues": str(error), + "runtime_error": str(error), + } + best_case_match = re.search( + r"([0-9.]+)%\s*>\s*([0-9.]+)%", str(error) + ) + if best_case_match: + row["reported_best_case_constraint_error_pct"] = float( + best_case_match.group(1) + ) + row["reported_allowed_constraint_error_pct"] = float( + best_case_match.group(2) + ) + rows.append(row) + del sim + gc.collect() + continue + + audit = build_calibration_audit( + X=X_current, + y_target=y_target, + weights=weights, + baseline_weights=baseline_weights, + calibration_event=calibration_event, + ss_values=ss_values, + ss_target=ss_target, + payroll_values=payroll_values, + payroll_target=payroll_target, + h6_income_values=None, + h6_revenue_target=None, + oasdi_tob_values=oasdi_tob_values if profile.use_tob else None, + oasdi_tob_target=load_oasdi_tob_projections(year) if profile.use_tob else None, + hi_tob_values=hi_tob_values if profile.use_tob else None, + hi_tob_target=load_hi_tob_projections(year) if profile.use_tob else None, + ) + audit["calibration_quality"] = classify_calibration_quality( + audit, + profile, + year=year, + ) + audit["age_bucket_size"] = age_bucket_size + audit["age_bucket_count"] = int(X_current.shape[1]) + + validation_issues = validate_calibration_audit( + audit, + profile, + year=year, + ) + audit["validation_issues"] = validation_issues + audit["validation_passed"] = not bool(validation_issues) + + row: dict[str, object] = { + "year": year, + "target_source": target_source, + "profile": profile.name, + "calibration_quality": audit["calibration_quality"], + "approximation_method": audit.get("approximation_method") + or audit.get("method_used"), + "iterations": iterations, + "age_bucket_size": age_bucket_size, + "window_max_constraint_error_pct": ( + approximate_window.max_constraint_error_pct + if approximate_window is not None + else profile.max_constraint_error_pct + ), + "window_max_age_error_pct": ( + approximate_window.max_age_error_pct + if approximate_window is not None + else profile.max_age_error_pct + ), + "max_constraint_pct_error": audit.get("max_constraint_pct_error"), + "age_max_pct_error": audit.get("age_max_pct_error"), + "positive_weight_count": audit.get("positive_weight_count"), + "effective_sample_size": audit.get("effective_sample_size"), + "top_10_weight_share_pct": audit.get("top_10_weight_share_pct"), + "top_100_weight_share_pct": audit.get("top_100_weight_share_pct"), + "negative_weight_pct": audit.get("negative_weight_pct"), + "validation_passed": audit["validation_passed"], + "validation_issue_count": len(validation_issues), + "validation_issues": "; ".join(validation_issues), + } + + tob_benchmarks = benchmark_tob_values( + year, + weights, + oasdi_tob_values=oasdi_tob_values, + hi_tob_values=hi_tob_values, + ) + if tob_benchmarks is not None: + row.update(tob_benchmarks) + + rows.append(row) + + del sim + gc.collect() + + return rows + + +def write_rows(rows: list[dict[str, object]], output: Path | None) -> None: + if not rows: + raise SystemExit("No rows to write.") + + fieldnames = [ + "year", + "target_source", + "profile", + "calibration_quality", + "approximation_method", + "iterations", + "age_bucket_size", + "window_max_constraint_error_pct", + "window_max_age_error_pct", + "max_constraint_pct_error", + "age_max_pct_error", + "positive_weight_count", + "effective_sample_size", + "top_10_weight_share_pct", + "top_100_weight_share_pct", + "negative_weight_pct", + "validation_passed", + "validation_issue_count", + "validation_issues", + "runtime_error", + "reported_best_case_constraint_error_pct", + "reported_allowed_constraint_error_pct", + "oasdi_tob_benchmark_pct_error", + "hi_tob_benchmark_pct_error", + ] + + if output is None: + writer = csv.DictWriter(sys.stdout, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + return + + output.parent.mkdir(parents=True, exist_ok=True) + with output.open("w", encoding="utf-8", newline="") as file: + writer = csv.DictWriter(file, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + + +def main() -> int: + args = parse_args() + rows = assess_years( + years=parse_years(args.years), + profile_name=args.profile, + target_source=args.target_source, + base_dataset_path=args.base_dataset, + ) + write_rows(rows, args.output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From ff099fd5b9b8cecaa4a5b1c0cf6bc86d650189a3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 31 Mar 2026 19:58:10 -0400 Subject: [PATCH 18/31] Add mixed-age household support diagnostic --- .../cps/long_term/support_augmentation.py | 217 +++++++++++++++++- .../test_long_term_calibration_contract.py | 47 ++++ 2 files changed, 263 insertions(+), 1 deletion(-) diff --git a/policyengine_us_data/datasets/cps/long_term/support_augmentation.py b/policyengine_us_data/datasets/cps/long_term/support_augmentation.py index 5ce817cae..8976a2424 100644 --- a/policyengine_us_data/datasets/cps/long_term/support_augmentation.py +++ b/policyengine_us_data/datasets/cps/long_term/support_augmentation.py @@ -79,12 +79,31 @@ class SinglePersonSyntheticGridRule: clone_weight_scale: float = 0.1 +@dataclass(frozen=True) +class MixedAgeAppendRule: + name: str + recipient_min_max_age: int + recipient_max_max_age: int + donor_min_max_age: int + donor_max_max_age: int + recipient_ss_state: ConstraintState = "positive" + recipient_payroll_state: ConstraintState = "any" + donor_ss_state: ConstraintState = "nonpositive" + donor_payroll_state: ConstraintState = "positive" + donor_age_shift: int = 0 + clone_weight_scale: float = 0.15 + + @dataclass(frozen=True) class SupportAugmentationProfile: name: str description: str rules: tuple[ - AgeShiftCloneRule | CompositePayrollRule | SinglePersonSyntheticGridRule, ... + AgeShiftCloneRule + | CompositePayrollRule + | SinglePersonSyntheticGridRule + | MixedAgeAppendRule, + ... ] @@ -353,6 +372,41 @@ class SupportAugmentationProfile: ), ) +LATE_MIXED_HOUSEHOLD_V1 = SupportAugmentationProfile( + name="late-mixed-household-v1", + description=( + "Append synthetic mixed-age households by combining older beneficiary " + "households with a younger payroll-rich donor person as a separate " + "subunit in the same household." + ), + rules=( + MixedAgeAppendRule( + name="older_75_84_plus_prime_age_earner", + recipient_min_max_age=75, + recipient_max_max_age=84, + donor_min_max_age=35, + donor_max_max_age=64, + recipient_ss_state="positive", + recipient_payroll_state="any", + donor_ss_state="nonpositive", + donor_payroll_state="positive", + clone_weight_scale=0.12, + ), + MixedAgeAppendRule( + name="older_85_plus_prime_age_earner", + recipient_min_max_age=85, + recipient_max_max_age=85, + donor_min_max_age=35, + donor_max_max_age=64, + recipient_ss_state="positive", + recipient_payroll_state="any", + donor_ss_state="nonpositive", + donor_payroll_state="positive", + clone_weight_scale=0.08, + ), + ), +) + NAMED_SUPPORT_AUGMENTATION_PROFILES = { LATE_CLONE_V1.name: LATE_CLONE_V1, @@ -361,6 +415,7 @@ class SupportAugmentationProfile: LATE_COMPOSITE_V2.name: LATE_COMPOSITE_V2, LATE_SYNTHETIC_GRID_V1.name: LATE_SYNTHETIC_GRID_V1, LATE_SYNTHETIC_GRID_V2.name: LATE_SYNTHETIC_GRID_V2, + LATE_MIXED_HOUSEHOLD_V1.name: LATE_MIXED_HOUSEHOLD_V1, } @@ -658,6 +713,47 @@ def _select_payroll_target_row( return adults[age_col].astype(float).idxmax() +def _clone_single_donor_person_row( + donor_row: pd.Series, + *, + base_year: int, + shared_household_id: int, + household_weight: float, + person_weight: float, + donor_age_shift: int, + id_counters: dict[str, int], +) -> tuple[pd.Series, dict[str, int]]: + cloned = donor_row.copy() + person_id_col = _period_column(PERSON_ID_COLUMN, base_year) + household_weight_col = _period_column("household_weight", base_year) + person_weight_col = _period_column("person_weight", base_year) + age_col = _period_column("age", base_year) + + cloned[person_id_col] = id_counters["person"] + id_counters["person"] += 1 + + for raw_column in ENTITY_ID_COLUMNS["household"]: + column = _period_column(raw_column, base_year) + cloned[column] = shared_household_id + + for entity_name, columns in ENTITY_ID_COLUMNS.items(): + if entity_name == "household": + continue + entity_id = id_counters[entity_name] + id_counters[entity_name] += 1 + for raw_column in columns: + column = _period_column(raw_column, base_year) + if column in cloned.index: + cloned[column] = entity_id + + cloned[age_col] = min(float(cloned[age_col]) + donor_age_shift, 85) + if household_weight_col in cloned.index: + cloned[household_weight_col] = household_weight + if person_weight_col in cloned.index: + cloned[person_weight_col] = person_weight + return cloned, id_counters + + def synthesize_composite_households( input_df: pd.DataFrame, *, @@ -898,6 +994,113 @@ def synthesize_single_person_grid_households( ) +def synthesize_mixed_age_households( + input_df: pd.DataFrame, + *, + base_year: int, + summary: pd.DataFrame, + rule: MixedAgeAppendRule, + id_counters: dict[str, int] | None = None, +) -> tuple[pd.DataFrame, dict[str, int], dict[str, Any]]: + recipient_ids = select_households_for_composite_rule( + summary, + min_max_age=rule.recipient_min_max_age, + max_max_age=rule.recipient_max_max_age, + ss_state=rule.recipient_ss_state, + payroll_state=rule.recipient_payroll_state, + ) + donor_ids = select_households_for_composite_rule( + summary, + min_max_age=rule.donor_min_max_age, + max_max_age=rule.donor_max_max_age, + ss_state=rule.donor_ss_state, + payroll_state=rule.donor_payroll_state, + ) + recipient_pairs = _quantile_pair_households(recipient_ids, donor_ids, summary) + if not recipient_pairs: + return ( + input_df.iloc[0:0].copy(), + id_counters.copy() if id_counters is not None else {}, + { + "rule": rule.name, + "recipient_household_count": int(len(recipient_ids)), + "donor_household_count": int(len(donor_ids)), + "synthetic_household_count": 0, + "synthetic_person_count": 0, + }, + ) + + recipient_clone_df, next_ids = clone_households_with_age_shift( + input_df, + base_year=base_year, + household_ids=pd.Index([recipient for recipient, _ in recipient_pairs]), + age_shift=0, + clone_weight_scale=rule.clone_weight_scale, + id_counters=id_counters, + ) + + household_id_col = _period_column("household_id", base_year) + person_weight_col = _period_column("person_weight", base_year) + + original_recipients = pd.unique( + input_df[ + input_df[household_id_col].isin( + [recipient for recipient, _ in recipient_pairs] + ) + ][household_id_col] + ) + cloned_household_ids = pd.unique(recipient_clone_df[household_id_col]) + cloned_mapping = { + int(original): int(cloned) + for original, cloned in zip(original_recipients, cloned_household_ids) + } + + donor_person_rows: list[pd.Series] = [] + for recipient_household_id, donor_household_id in recipient_pairs: + cloned_household_id = cloned_mapping[int(recipient_household_id)] + cloned_household_rows = recipient_clone_df[ + recipient_clone_df[household_id_col] == cloned_household_id + ] + donor_rows = input_df[input_df[household_id_col] == int(donor_household_id)] + donor_row_idx = _select_payroll_target_row( + donor_rows, + base_year=base_year, + ) + household_weight = float(cloned_household_rows.iloc[0][_period_column("household_weight", base_year)]) + person_weight = ( + float(cloned_household_rows.iloc[0][person_weight_col]) + if person_weight_col in cloned_household_rows.columns + else household_weight + ) + donor_row, next_ids = _clone_single_donor_person_row( + donor_rows.loc[donor_row_idx], + base_year=base_year, + shared_household_id=cloned_household_id, + household_weight=household_weight, + person_weight=person_weight, + donor_age_shift=rule.donor_age_shift, + id_counters=next_ids, + ) + donor_person_rows.append(donor_row) + + donor_df = pd.DataFrame(donor_person_rows, columns=input_df.columns) + synthetic_df = pd.concat([recipient_clone_df, donor_df], ignore_index=True) + return ( + synthetic_df, + next_ids, + { + "rule": rule.name, + "recipient_household_count": int(len(recipient_ids)), + "donor_household_count": int(len(donor_ids)), + "synthetic_household_count": int( + synthetic_df[household_id_col].nunique() + ), + "synthetic_person_count": int(len(synthetic_df)), + "donor_age_shift": rule.donor_age_shift, + }, + ) + + def augment_input_dataframe( input_df: pd.DataFrame, *, @@ -960,6 +1163,18 @@ def augment_input_dataframe( rule_reports.append(composite_report) continue + if isinstance(rule, MixedAgeAppendRule): + mixed_df, id_counters, mixed_report = synthesize_mixed_age_households( + input_df, + base_year=base_year, + summary=summary, + rule=rule, + id_counters=id_counters, + ) + clone_frames.append(mixed_df) + rule_reports.append(mixed_report) + continue + synthetic_df, id_counters, synthetic_report = synthesize_single_person_grid_households( input_df, base_year=base_year, diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index e6e31de83..23a9c4671 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -40,6 +40,7 @@ from policyengine_us_data.datasets.cps.long_term.support_augmentation import ( AgeShiftCloneRule, CompositePayrollRule, + MixedAgeAppendRule, SinglePersonSyntheticGridRule, SupportAugmentationProfile, augment_input_dataframe, @@ -292,6 +293,52 @@ def test_support_augmentation_appends_single_person_synthetic_grid_households(): ) == {50_000.0} +def test_support_augmentation_appends_mixed_age_household(): + import pandas as pd + + df = pd.DataFrame(_toy_support_dataframe()) + profile = SupportAugmentationProfile( + name="mixed-age-profile", + description="Toy mixed-age household support augmentation profile.", + rules=( + MixedAgeAppendRule( + name="older_plus_younger_earner", + recipient_min_max_age=75, + recipient_max_max_age=84, + donor_min_max_age=55, + donor_max_max_age=64, + recipient_ss_state="positive", + recipient_payroll_state="any", + donor_ss_state="nonpositive", + donor_payroll_state="positive", + clone_weight_scale=0.2, + ), + ), + ) + augmented_df, report = augment_input_dataframe( + df, + base_year=2024, + profile=profile, + ) + assert report["base_household_count"] == 3 + assert report["augmented_household_count"] == 4 + synthetic_household_ids = set( + augmented_df["household_id__2024"].unique() + ) - {1, 2, 3} + assert len(synthetic_household_ids) == 1 + synthetic_rows = augmented_df[ + augmented_df["household_id__2024"].isin(synthetic_household_ids) + ] + assert sorted(synthetic_rows["age__2024"].tolist()) == [60.0, 77.0, 80.0] + assert synthetic_rows["social_security_retirement__2024"].sum() == pytest.approx( + 30_000.0 + ) + assert synthetic_rows["employment_income_before_lsr__2024"].sum() == pytest.approx( + 62_000.0 + ) + assert synthetic_rows["tax_unit_id__2024"].nunique() == 3 + + def test_age_bin_helpers_preserve_population_totals(): bins = build_age_bins(n_ages=86, bucket_size=5) assert bins[0] == (0, 5) From 7a03b8b1a59b0788638ce3cf3bb6d0db7aaa1d6a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 31 Mar 2026 22:39:38 -0400 Subject: [PATCH 19/31] Prototype minimal synthetic support for 2100 --- .../prototype_synthetic_2100_support.py | 820 ++++++++++++++++++ 1 file changed, 820 insertions(+) create mode 100644 policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py diff --git a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py new file mode 100644 index 000000000..86faf86b2 --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py @@ -0,0 +1,820 @@ +from __future__ import annotations + +import argparse +from dataclasses import asdict, dataclass +import json +from pathlib import Path + +import numpy as np +import pandas as pd +from policyengine_us import Microsimulation + +from calibration import assess_nonnegative_feasibility, calibrate_entropy +from projection_utils import aggregate_age_targets, build_age_bins +from ssa_data import ( + get_long_term_target_source, + load_ssa_age_projections, + load_ssa_benefit_projections, + load_taxable_payroll_projections, + set_long_term_target_source, +) + + +DEFAULT_DATASET = "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5" +DEFAULT_YEAR = 2100 +BASE_YEAR = 2024 + + +@dataclass(frozen=True) +class SyntheticTemplate: + name: str + head_ages: tuple[int, ...] + spouse_age_offsets: tuple[int | None, ...] + dependent_age_sets: tuple[tuple[int, ...], ...] + ss_source: str + payroll_source: str + pension_source: str + dividend_source: str + ss_split: tuple[float, float] + payroll_split: tuple[float, float] + + +@dataclass(frozen=True) +class SyntheticCandidate: + archetype: str + head_age: int + spouse_age: int | None + dependent_ages: tuple[int, ...] + head_wages: float + spouse_wages: float + head_ss: float + spouse_ss: float + pension_income: float + dividend_income: float + + @property + def payroll_total(self) -> float: + return float(self.head_wages + self.spouse_wages) + + @property + def ss_total(self) -> float: + return float(self.head_ss + self.spouse_ss) + + def ages(self) -> list[int]: + values = [self.head_age] + if self.spouse_age is not None: + values.append(self.spouse_age) + values.extend(self.dependent_ages) + return values + + def filing_status(self) -> str: + return "joint" if self.spouse_age is not None else "single" + + def taxable_benefits_proxy(self) -> float: + benefits = self.ss_total + if benefits <= 0: + return 0.0 + provisional_income = ( + self.payroll_total + + self.pension_income + + self.dividend_income + + 0.5 * benefits + ) + if self.filing_status() == "joint": + base = 32_000.0 + adjusted = 44_000.0 + lesser_cap = 6_000.0 + else: + base = 25_000.0 + adjusted = 34_000.0 + lesser_cap = 4_500.0 + + if provisional_income <= base: + return 0.0 + if provisional_income <= adjusted: + return min(0.5 * benefits, 0.5 * (provisional_income - base)) + return min( + 0.85 * benefits, + 0.85 * (provisional_income - adjusted) + min(0.5 * benefits, lesser_cap), + ) + + +TEMPLATES = ( + SyntheticTemplate( + name="older_beneficiary_single", + head_ages=(67, 75, 85), + spouse_age_offsets=(None,), + dependent_age_sets=((),), + ss_source="older_beneficiary", + payroll_source="zero", + pension_source="older_asset", + dividend_source="older_asset", + ss_split=(1.0, 0.0), + payroll_split=(0.0, 0.0), + ), + SyntheticTemplate( + name="older_beneficiary_couple", + head_ages=(67, 75, 85), + spouse_age_offsets=(-2, -5), + dependent_age_sets=((),), + ss_source="older_couple_beneficiary", + payroll_source="zero", + pension_source="older_asset", + dividend_source="older_asset", + ss_split=(0.55, 0.45), + payroll_split=(0.0, 0.0), + ), + SyntheticTemplate( + name="older_worker_single", + head_ages=(67, 75, 80), + spouse_age_offsets=(None,), + dependent_age_sets=((),), + ss_source="older_worker", + payroll_source="older_worker", + pension_source="older_asset", + dividend_source="older_asset", + ss_split=(1.0, 0.0), + payroll_split=(1.0, 0.0), + ), + SyntheticTemplate( + name="older_worker_couple", + head_ages=(67, 75, 80), + spouse_age_offsets=(-2,), + dependent_age_sets=((),), + ss_source="older_worker", + payroll_source="older_worker", + pension_source="older_asset", + dividend_source="older_asset", + ss_split=(0.55, 0.45), + payroll_split=(0.55, 0.45), + ), + SyntheticTemplate( + name="mixed_retiree_worker_couple", + head_ages=(67, 75, 85), + spouse_age_offsets=(-20, -25, -35), + dependent_age_sets=((),), + ss_source="older_beneficiary", + payroll_source="prime_worker", + pension_source="older_asset", + dividend_source="older_asset", + ss_split=(1.0, 0.0), + payroll_split=(0.0, 1.0), + ), + SyntheticTemplate( + name="prime_worker_single", + head_ages=(22, 27, 35, 45, 60), + spouse_age_offsets=(None,), + dependent_age_sets=((),), + ss_source="zero", + payroll_source="prime_worker", + pension_source="prime_asset", + dividend_source="prime_asset", + ss_split=(0.0, 0.0), + payroll_split=(1.0, 0.0), + ), + SyntheticTemplate( + name="prime_worker_couple", + head_ages=(27, 40, 55), + spouse_age_offsets=(-2,), + dependent_age_sets=((),), + ss_source="zero", + payroll_source="prime_worker", + pension_source="prime_asset", + dividend_source="prime_asset", + ss_split=(0.0, 0.0), + payroll_split=(0.6, 0.4), + ), + SyntheticTemplate( + name="prime_worker_family", + head_ages=(27, 40, 55), + spouse_age_offsets=(-2,), + dependent_age_sets=((1, 3), (6, 11), (15, 17)), + ss_source="zero", + payroll_source="prime_worker_family", + pension_source="prime_asset", + dividend_source="prime_asset", + ss_split=(0.0, 0.0), + payroll_split=(0.6, 0.4), + ), + SyntheticTemplate( + name="older_plus_prime_worker_family", + head_ages=(67, 75, 85), + spouse_age_offsets=(-25, -35), + dependent_age_sets=((1, 3), (10, 15)), + ss_source="older_beneficiary", + payroll_source="prime_worker_family", + pension_source="older_asset", + dividend_source="older_asset", + ss_split=(1.0, 0.0), + payroll_split=(0.0, 1.0), + ), +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Prototype a fully synthetic late-year support using a minimal set of " + "head/spouse/dependent features." + ) + ) + parser.add_argument( + "--year", + type=int, + default=DEFAULT_YEAR, + help="Projection year to target.", + ) + parser.add_argument( + "--target-source", + default=get_long_term_target_source(), + help="Named long-run target source package.", + ) + parser.add_argument( + "--base-dataset", + default=DEFAULT_DATASET, + help="Base 2024 dataset used to derive comparison pools.", + ) + parser.add_argument( + "--output", + type=Path, + help="Optional JSON output path.", + ) + return parser.parse_args() + + +def classify_archetype( + *, + head_age: float, + spouse_age: float | None, + dependent_count: int, + ss_total: float, + payroll_total: float, +) -> str: + older_head = head_age >= 65 + has_spouse = spouse_age is not None + older_spouse = spouse_age is not None and spouse_age >= 65 + positive_ss = ss_total > 0 + positive_payroll = payroll_total > 0 + + if older_head: + if dependent_count > 0 and positive_payroll: + return "older_plus_prime_worker_family" + if has_spouse and spouse_age is not None and spouse_age < 65 and positive_payroll: + return "mixed_retiree_worker_couple" + if has_spouse and older_spouse: + if positive_payroll and positive_ss: + return "older_worker_couple" + if positive_ss: + return "older_beneficiary_couple" + return "older_couple_other" + if positive_payroll and positive_ss: + return "older_worker_single" + if positive_ss: + return "older_beneficiary_single" + return "older_single_other" + + if dependent_count > 0: + return "prime_worker_family" if positive_payroll else "prime_other_family" + if has_spouse: + return "prime_worker_couple" if positive_payroll else "prime_other_couple" + return "prime_worker_single" if positive_payroll else "prime_other_single" + + +def build_actual_tax_unit_summary(base_dataset: str) -> pd.DataFrame: + sim = Microsimulation(dataset=base_dataset) + input_df = sim.to_input_dataframe() + + person_df = pd.DataFrame( + { + "tax_unit_id": sim.calculate("person_tax_unit_id", period=BASE_YEAR).values, + "age": sim.calculate("age", period=BASE_YEAR).values, + "is_head": sim.calculate("is_tax_unit_head", period=BASE_YEAR).values, + "is_spouse": sim.calculate("is_tax_unit_spouse", period=BASE_YEAR).values, + "is_dependent": sim.calculate( + "is_tax_unit_dependent", period=BASE_YEAR + ).values, + "social_security": sim.calculate( + "social_security", period=BASE_YEAR + ).values, + "payroll": ( + sim.calculate("employment_income_before_lsr", period=BASE_YEAR).values + + sim.calculate( + "self_employment_income_before_lsr", period=BASE_YEAR + ).values + ), + "dividend_income": sim.calculate( + "qualified_dividend_income", period=BASE_YEAR + ).values, + "pension_income": sim.calculate( + "taxable_pension_income", period=BASE_YEAR + ).values, + "person_weight": input_df[f"person_weight__{BASE_YEAR}"].astype(float).values, + } + ) + + rows: list[dict[str, object]] = [] + for tax_unit_id, group in person_df.groupby("tax_unit_id", sort=False): + heads = group[group["is_head"]] + spouses = group[group["is_spouse"]] + dependents = group[group["is_dependent"]] + adults = group[group["age"] >= 18] + head_age = float( + heads["age"].iloc[0] + if not heads.empty + else adults["age"].max() if not adults.empty else group["age"].max() + ) + spouse_age = ( + float(spouses["age"].iloc[0]) if not spouses.empty else None + ) + dependent_count = int(len(dependents)) + row = { + "tax_unit_id": int(tax_unit_id), + "head_age": head_age, + "spouse_age": spouse_age, + "dependent_count": dependent_count, + "payroll_total": float(group["payroll"].sum()), + "ss_total": float(group["social_security"].sum()), + "dividend_income": float(group["dividend_income"].sum()), + "pension_income": float(group["pension_income"].sum()), + "support_count_weight": 1.0, + "person_weight_proxy": float(group["person_weight"].max()), + } + row["archetype"] = classify_archetype( + head_age=row["head_age"], + spouse_age=row["spouse_age"], + dependent_count=row["dependent_count"], + ss_total=row["ss_total"], + payroll_total=row["payroll_total"], + ) + rows.append(row) + + return pd.DataFrame(rows) + + +def load_base_aggregates(base_dataset: str) -> dict[str, float]: + sim = Microsimulation(dataset=base_dataset) + household_series = sim.calculate("household_id", period=BASE_YEAR, map_to="household") + weights = household_series.weights.values.astype(float) + ss = sim.calculate("social_security", period=BASE_YEAR, map_to="household").values + payroll = ( + sim.calculate( + "taxable_earnings_for_social_security", + period=BASE_YEAR, + map_to="household", + ).values + + sim.calculate( + "social_security_taxable_self_employment_income", + period=BASE_YEAR, + map_to="household", + ).values + ) + return { + "weighted_ss_total": float(np.sum(ss * weights)), + "weighted_payroll_total": float(np.sum(payroll * weights)), + } + + +def quantile_levels( + values: pd.Series, + *, + quantiles: tuple[float, ...], + include_zero: bool = False, + positive_only: bool = False, +) -> list[float]: + series = values.astype(float) + if positive_only: + series = series[series > 0] + if series.empty: + levels = [0.0] + else: + levels = [float(series.quantile(q)) for q in quantiles] + if include_zero: + levels = [0.0, *levels] + deduped = [] + for value in levels: + rounded = round(value, 2) + if rounded not in deduped: + deduped.append(rounded) + return [float(value) for value in deduped] + + +def _scale_levels(levels: list[float], scale: float) -> list[float]: + return [round(level * scale, 2) for level in levels] + + +def build_quantile_pools( + actual_summary: pd.DataFrame, + *, + ss_scale: float, + earnings_scale: float, +) -> dict[str, dict[str, list[float]]]: + masks = { + "older_beneficiary": (actual_summary["head_age"] >= 65) + & (actual_summary["ss_total"] > 0), + "older_couple_beneficiary": (actual_summary["head_age"] >= 65) + & actual_summary["spouse_age"].fillna(-1).ge(65) + & (actual_summary["ss_total"] > 0), + "older_worker": (actual_summary["head_age"] >= 65) + & (actual_summary["ss_total"] > 0) + & (actual_summary["payroll_total"] > 0), + "prime_worker": (actual_summary["head_age"] < 65) + & (actual_summary["payroll_total"] > 0), + "prime_worker_family": (actual_summary["head_age"] < 65) + & (actual_summary["dependent_count"] > 0) + & (actual_summary["payroll_total"] > 0), + "older_asset": (actual_summary["head_age"] >= 65) + & ( + (actual_summary["pension_income"] > 0) + | (actual_summary["dividend_income"] > 0) + ), + "prime_asset": (actual_summary["head_age"] < 65) + & ( + (actual_summary["pension_income"] > 0) + | (actual_summary["dividend_income"] > 0) + ), + "zero": actual_summary["head_age"].notna(), + } + + pools: dict[str, dict[str, list[float]]] = {} + for name, mask in masks.items(): + subset = actual_summary[mask] + pools[name] = { + "ss": _scale_levels( + quantile_levels( + subset["ss_total"], + quantiles=(0.25, 0.5, 0.75), + include_zero=(name == "zero"), + positive_only=(name != "zero"), + ), + ss_scale, + ), + "payroll": _scale_levels( + quantile_levels( + subset["payroll_total"], + quantiles=(0.25, 0.5, 0.75), + include_zero=( + name + in { + "zero", + "older_beneficiary", + "older_couple_beneficiary", + "older_asset", + } + ), + positive_only=(name not in {"zero"}), + ), + earnings_scale, + ), + "pension": _scale_levels( + quantile_levels( + subset["pension_income"], + quantiles=(0.5, 0.9), + include_zero=True, + positive_only=True, + ), + earnings_scale, + ), + "dividend": _scale_levels( + quantile_levels( + subset["dividend_income"], + quantiles=(0.5, 0.9), + include_zero=True, + positive_only=True, + ), + earnings_scale, + ), + } + return pools + + +def generate_synthetic_candidates( + pools: dict[str, dict[str, list[float]]], +) -> list[SyntheticCandidate]: + candidates: list[SyntheticCandidate] = [] + for template in TEMPLATES: + ss_levels = pools[template.ss_source]["ss"] + payroll_levels = pools[template.payroll_source]["payroll"] + pension_levels = pools[template.pension_source]["pension"] + dividend_levels = pools[template.dividend_source]["dividend"] + for head_age in template.head_ages: + for spouse_offset in template.spouse_age_offsets: + spouse_age = ( + None if spouse_offset is None else max(18, head_age + spouse_offset) + ) + for dependent_ages in template.dependent_age_sets: + for ss_total in ss_levels: + for payroll_total in payroll_levels: + for pension_income in pension_levels: + for dividend_income in dividend_levels: + head_ss = ss_total * template.ss_split[0] + spouse_ss = ss_total * template.ss_split[1] + head_wages = payroll_total * template.payroll_split[0] + spouse_wages = payroll_total * template.payroll_split[1] + candidates.append( + SyntheticCandidate( + archetype=template.name, + head_age=head_age, + spouse_age=spouse_age, + dependent_ages=tuple(dependent_ages), + head_wages=head_wages, + spouse_wages=spouse_wages, + head_ss=head_ss, + spouse_ss=spouse_ss, + pension_income=pension_income, + dividend_income=dividend_income, + ) + ) + # Deduplicate exact duplicates caused by repeated quantiles. + deduped: dict[tuple[object, ...], SyntheticCandidate] = {} + for candidate in candidates: + key = ( + candidate.archetype, + candidate.head_age, + candidate.spouse_age, + candidate.dependent_ages, + round(candidate.head_wages, 2), + round(candidate.spouse_wages, 2), + round(candidate.head_ss, 2), + round(candidate.spouse_ss, 2), + round(candidate.pension_income, 2), + round(candidate.dividend_income, 2), + ) + deduped[key] = candidate + return list(deduped.values()) + + +def age_bucket_vector(ages: list[int], age_bins: list[tuple[int, int]]) -> np.ndarray: + vector = np.zeros(len(age_bins), dtype=float) + for age in ages: + if age >= 85: + vector[-1] += 1.0 + continue + for idx, (start, end) in enumerate(age_bins): + if start <= age < end: + vector[idx] += 1.0 + break + return vector + + +def solve_synthetic_support( + candidates: list[SyntheticCandidate], + *, + year: int, +) -> tuple[np.ndarray, dict[str, object]]: + age_targets = load_ssa_age_projections(start_year=year, end_year=year) + age_bins = build_age_bins(n_ages=age_targets.shape[0], bucket_size=5) + aggregated_age_targets = aggregate_age_targets(age_targets, age_bins)[:, 0] + + X = np.vstack([age_bucket_vector(candidate.ages(), age_bins) for candidate in candidates]) + ss_values = np.array([candidate.ss_total for candidate in candidates], dtype=float) + payroll_values = np.array( + [candidate.payroll_total for candidate in candidates], dtype=float + ) + baseline_weights = np.ones(len(candidates), dtype=float) + + try: + weights, iterations = calibrate_entropy( + X, + aggregated_age_targets, + baseline_weights, + ss_values=ss_values, + ss_target=load_ssa_benefit_projections(year), + payroll_values=payroll_values, + payroll_target=load_taxable_payroll_projections(year), + n_ages=X.shape[1], + max_iters=500, + tol=1e-9, + ) + return weights, { + "method": "entropy", + "iterations": int(iterations), + "best_case_max_pct_error": 0.0, + "age_bucket_size": 5, + } + except RuntimeError as error: + feasibility = assess_nonnegative_feasibility( + np.column_stack([X, ss_values, payroll_values]), + np.concatenate( + [ + aggregated_age_targets, + np.array( + [ + load_ssa_benefit_projections(year), + load_taxable_payroll_projections(year), + ] + ), + ] + ), + return_weights=True, + ) + if not feasibility["success"] or feasibility.get("weights") is None: + raise RuntimeError( + f"Synthetic support could not match {year} targets: {error}" + ) from error + return np.asarray(feasibility["weights"], dtype=float), { + "method": "lp_minimax", + "iterations": 1, + "best_case_max_pct_error": feasibility["best_case_max_pct_error"], + "age_bucket_size": 5, + "entropy_error": str(error), + } + + +def summarize_solution( + candidates: list[SyntheticCandidate], + weights: np.ndarray, + actual_summary: pd.DataFrame, +) -> dict[str, object]: + weight_sum = float(weights.sum()) + candidate_df = pd.DataFrame( + [ + { + **asdict(candidate), + "spouse_age": candidate.spouse_age, + "dependent_count": len(candidate.dependent_ages), + "payroll_total": candidate.payroll_total, + "ss_total": candidate.ss_total, + "taxable_benefits_proxy": candidate.taxable_benefits_proxy(), + "synthetic_weight": float(weight), + } + for candidate, weight in zip(candidates, weights) + ] + ) + candidate_df["weight_share_pct"] = ( + candidate_df["synthetic_weight"] / weight_sum * 100 if weight_sum > 0 else 0.0 + ) + candidate_df = candidate_df.sort_values("synthetic_weight", ascending=False) + + def _weighted_mean(group: pd.DataFrame, column: str) -> float: + total = float(group["synthetic_weight"].sum()) + if total <= 0: + return 0.0 + return float(np.average(group[column], weights=group["synthetic_weight"])) + + synthetic_rows = [] + for archetype, group in candidate_df.groupby("archetype", sort=False): + synthetic_rows.append( + { + "archetype": archetype, + "synthetic_weight": float(group["synthetic_weight"].sum()), + "candidate_count": int(len(group)), + "avg_head_age": _weighted_mean(group, "head_age"), + "avg_payroll_total": _weighted_mean(group, "payroll_total"), + "avg_ss_total": _weighted_mean(group, "ss_total"), + "avg_pension_income": _weighted_mean(group, "pension_income"), + "avg_dividend_income": _weighted_mean(group, "dividend_income"), + } + ) + synthetic_archetypes = pd.DataFrame(synthetic_rows).sort_values( + "synthetic_weight", + ascending=False, + ) + synthetic_archetypes["synthetic_weight_share_pct"] = ( + synthetic_archetypes["synthetic_weight"] / weight_sum * 100 + if weight_sum > 0 + else 0.0 + ) + + actual_archetypes = ( + actual_summary.groupby("archetype", as_index=False) + .agg( + actual_support_count=("archetype", "count"), + avg_head_age=("head_age", "mean"), + avg_payroll_total=("payroll_total", "mean"), + avg_ss_total=("ss_total", "mean"), + avg_pension_income=("pension_income", "mean"), + avg_dividend_income=("dividend_income", "mean"), + ) + .sort_values("actual_support_count", ascending=False) + ) + actual_total = float(actual_archetypes["actual_support_count"].sum()) + actual_archetypes["actual_support_share_pct"] = ( + actual_archetypes["actual_support_count"] / actual_total * 100 + if actual_total > 0 + else 0.0 + ) + + comparison = pd.merge( + synthetic_archetypes[ + ["archetype", "synthetic_weight_share_pct", "candidate_count"] + ], + actual_archetypes[ + ["archetype", "actual_support_share_pct", "actual_support_count"] + ], + on="archetype", + how="outer", + ).fillna(0.0) + comparison["share_gap_pct_points"] = ( + comparison["synthetic_weight_share_pct"] + - comparison["actual_support_share_pct"] + ) + comparison = comparison.sort_values( + "synthetic_weight_share_pct", + ascending=False, + ) + + weighted_metrics = { + "synthetic_payroll_positive_85_plus_household_share_pct": float( + candidate_df.loc[ + (candidate_df["head_age"] >= 85) & (candidate_df["payroll_total"] > 0), + "synthetic_weight", + ].sum() + / weight_sum + * 100 + ) + if weight_sum > 0 + else 0.0, + "synthetic_mixed_retiree_worker_share_pct": float( + candidate_df.loc[ + candidate_df["archetype"].eq("mixed_retiree_worker_couple"), + "synthetic_weight", + ].sum() + / weight_sum + * 100 + ) + if weight_sum > 0 + else 0.0, + "synthetic_units_with_positive_pension_or_dividend_share_pct": float( + candidate_df.loc[ + (candidate_df["pension_income"] > 0) + | (candidate_df["dividend_income"] > 0), + "synthetic_weight", + ].sum() + / weight_sum + * 100 + ) + if weight_sum > 0 + else 0.0, + "synthetic_avg_taxable_benefits_proxy_share_pct": float( + ( + (candidate_df["taxable_benefits_proxy"] * candidate_df["synthetic_weight"]).sum() + / max((candidate_df["ss_total"] * candidate_df["synthetic_weight"]).sum(), 1.0) + ) + * 100 + ), + } + + return { + "synthetic_candidate_count": int(len(candidate_df)), + "positive_weight_candidate_count": int((candidate_df["synthetic_weight"] > 0).sum()), + "top_candidates": candidate_df.head(20).to_dict("records"), + "synthetic_archetypes": synthetic_archetypes.to_dict("records"), + "actual_support_archetypes": actual_archetypes.to_dict("records"), + "archetype_gap_table": comparison.to_dict("records"), + "weighted_metrics": weighted_metrics, + } + + +def main() -> int: + args = parse_args() + set_long_term_target_source(args.target_source) + + actual_summary = build_actual_tax_unit_summary(args.base_dataset) + base_aggregates = load_base_aggregates(args.base_dataset) + ss_scale = load_ssa_benefit_projections(args.year) / max( + base_aggregates["weighted_ss_total"], + 1.0, + ) + earnings_scale = load_taxable_payroll_projections(args.year) / max( + base_aggregates["weighted_payroll_total"], + 1.0, + ) + pools = build_quantile_pools( + actual_summary, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + ) + candidates = generate_synthetic_candidates(pools) + weights, solve_info = solve_synthetic_support(candidates, year=args.year) + solution_summary = summarize_solution(candidates, weights, actual_summary) + + report = { + "year": args.year, + "target_source": args.target_source, + "base_dataset": args.base_dataset, + "solve_info": solve_info, + "targets": { + "ss_total": float(load_ssa_benefit_projections(args.year)), + "taxable_payroll": float(load_taxable_payroll_projections(args.year)), + }, + "macro_scales": { + "ss_scale": float(ss_scale), + "earnings_scale": float(earnings_scale), + **base_aggregates, + }, + "actual_support_tax_unit_count": int(len(actual_summary)), + "synthetic_solution": solution_summary, + } + + payload = json.dumps(report, indent=2) + if args.output is None: + print(payload) + else: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(payload + "\n", encoding="utf-8") + print(args.output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From c14bbb884c086cb7d38fe42bf136e60dcf713319 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 1 Apr 2026 07:25:50 -0400 Subject: [PATCH 20/31] Expand synthetic 2100 support prototype --- .../prototype_synthetic_2100_support.py | 188 ++++++++++++++---- 1 file changed, 147 insertions(+), 41 deletions(-) diff --git a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py index 86faf86b2..001a4418c 100644 --- a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py +++ b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py @@ -37,6 +37,10 @@ class SyntheticTemplate: dividend_source: str ss_split: tuple[float, float] payroll_split: tuple[float, float] + ss_scale_factors: tuple[float, ...] = (1.0,) + payroll_scale_factors: tuple[float, ...] = (1.0,) + pension_scale_factors: tuple[float, ...] = (1.0,) + dividend_scale_factors: tuple[float, ...] = (1.0,) @dataclass(frozen=True) @@ -102,7 +106,7 @@ def taxable_benefits_proxy(self) -> float: TEMPLATES = ( SyntheticTemplate( name="older_beneficiary_single", - head_ages=(67, 75, 85), + head_ages=(62, 67, 72, 77, 82, 85), spouse_age_offsets=(None,), dependent_age_sets=((),), ss_source="older_beneficiary", @@ -111,11 +115,14 @@ def taxable_benefits_proxy(self) -> float: dividend_source="older_asset", ss_split=(1.0, 0.0), payroll_split=(0.0, 0.0), + ss_scale_factors=(0.75, 1.0, 1.25), + pension_scale_factors=(0.0, 1.0), + dividend_scale_factors=(0.0, 1.0), ), SyntheticTemplate( name="older_beneficiary_couple", - head_ages=(67, 75, 85), - spouse_age_offsets=(-2, -5), + head_ages=(62, 67, 72, 77, 82, 85), + spouse_age_offsets=(-2, -5, -8), dependent_age_sets=((),), ss_source="older_couple_beneficiary", payroll_source="zero", @@ -123,10 +130,13 @@ def taxable_benefits_proxy(self) -> float: dividend_source="older_asset", ss_split=(0.55, 0.45), payroll_split=(0.0, 0.0), + ss_scale_factors=(0.75, 1.0, 1.25), + pension_scale_factors=(0.0, 1.0), + dividend_scale_factors=(0.0, 1.0), ), SyntheticTemplate( name="older_worker_single", - head_ages=(67, 75, 80), + head_ages=(62, 65, 67, 70, 75, 80), spouse_age_offsets=(None,), dependent_age_sets=((),), ss_source="older_worker", @@ -135,11 +145,15 @@ def taxable_benefits_proxy(self) -> float: dividend_source="older_asset", ss_split=(1.0, 0.0), payroll_split=(1.0, 0.0), + ss_scale_factors=(0.5, 0.75, 1.0, 1.25), + payroll_scale_factors=(0.5, 1.0, 1.5, 2.0), + pension_scale_factors=(0.0, 1.0), + dividend_scale_factors=(0.0, 1.0), ), SyntheticTemplate( name="older_worker_couple", - head_ages=(67, 75, 80), - spouse_age_offsets=(-2,), + head_ages=(62, 65, 67, 70, 75, 80), + spouse_age_offsets=(-2, -5), dependent_age_sets=((),), ss_source="older_worker", payroll_source="older_worker", @@ -147,11 +161,15 @@ def taxable_benefits_proxy(self) -> float: dividend_source="older_asset", ss_split=(0.55, 0.45), payroll_split=(0.55, 0.45), + ss_scale_factors=(0.5, 0.75, 1.0, 1.25), + payroll_scale_factors=(0.5, 1.0, 1.5, 2.0), + pension_scale_factors=(0.0, 1.0), + dividend_scale_factors=(0.0, 1.0), ), SyntheticTemplate( name="mixed_retiree_worker_couple", - head_ages=(67, 75, 85), - spouse_age_offsets=(-20, -25, -35), + head_ages=(62, 67, 72, 77, 82, 85), + spouse_age_offsets=(-10, -15, -20, -25, -35), dependent_age_sets=((),), ss_source="older_beneficiary", payroll_source="prime_worker", @@ -159,10 +177,14 @@ def taxable_benefits_proxy(self) -> float: dividend_source="older_asset", ss_split=(1.0, 0.0), payroll_split=(0.0, 1.0), + ss_scale_factors=(0.5, 0.75, 1.0, 1.25), + payroll_scale_factors=(0.5, 1.0, 1.5, 2.0), + pension_scale_factors=(0.0, 1.0), + dividend_scale_factors=(0.0, 1.0), ), SyntheticTemplate( name="prime_worker_single", - head_ages=(22, 27, 35, 45, 60), + head_ages=(20, 22, 25, 27, 30, 35, 40, 45, 50, 55, 60, 64), spouse_age_offsets=(None,), dependent_age_sets=((),), ss_source="zero", @@ -171,11 +193,14 @@ def taxable_benefits_proxy(self) -> float: dividend_source="prime_asset", ss_split=(0.0, 0.0), payroll_split=(1.0, 0.0), + payroll_scale_factors=(0.5, 1.0, 1.5, 2.0), + pension_scale_factors=(0.0, 1.0), + dividend_scale_factors=(0.0, 1.0), ), SyntheticTemplate( name="prime_worker_couple", - head_ages=(27, 40, 55), - spouse_age_offsets=(-2,), + head_ages=(25, 30, 35, 40, 45, 50, 55, 60), + spouse_age_offsets=(-2, -5, -8), dependent_age_sets=((),), ss_source="zero", payroll_source="prime_worker", @@ -183,30 +208,72 @@ def taxable_benefits_proxy(self) -> float: dividend_source="prime_asset", ss_split=(0.0, 0.0), payroll_split=(0.6, 0.4), + payroll_scale_factors=(0.5, 1.0, 1.5, 2.0), + pension_scale_factors=(0.0, 1.0), + dividend_scale_factors=(0.0, 1.0), ), SyntheticTemplate( name="prime_worker_family", - head_ages=(27, 40, 55), + head_ages=(25, 30, 35, 40, 45, 50, 55), spouse_age_offsets=(-2,), - dependent_age_sets=((1, 3), (6, 11), (15, 17)), + dependent_age_sets=((0,), (3,), (7,), (12,), (16,), (4, 9), (11, 16)), ss_source="zero", payroll_source="prime_worker_family", pension_source="prime_asset", dividend_source="prime_asset", ss_split=(0.0, 0.0), payroll_split=(0.6, 0.4), + payroll_scale_factors=(0.5, 1.0, 1.5, 2.0), + pension_scale_factors=(0.0, 1.0), + dividend_scale_factors=(0.0, 1.0), ), SyntheticTemplate( name="older_plus_prime_worker_family", - head_ages=(67, 75, 85), - spouse_age_offsets=(-25, -35), - dependent_age_sets=((1, 3), (10, 15)), + head_ages=(62, 67, 72, 77, 82, 85), + spouse_age_offsets=(-15, -25, -35), + dependent_age_sets=((0,), (7,), (15,), (4, 9), (11, 16)), ss_source="older_beneficiary", payroll_source="prime_worker_family", pension_source="older_asset", dividend_source="older_asset", ss_split=(1.0, 0.0), payroll_split=(0.0, 1.0), + ss_scale_factors=(0.5, 0.75, 1.0, 1.25), + payroll_scale_factors=(0.5, 1.0, 1.5, 2.0), + pension_scale_factors=(0.0, 1.0), + dividend_scale_factors=(0.0, 1.0), + ), + SyntheticTemplate( + name="late_worker_couple", + head_ages=(58, 60, 62, 64, 66, 68), + spouse_age_offsets=(-2, -5), + dependent_age_sets=((),), + ss_source="older_worker", + payroll_source="prime_worker", + pension_source="prime_asset", + dividend_source="prime_asset", + ss_split=(0.6, 0.4), + payroll_split=(0.6, 0.4), + ss_scale_factors=(0.25, 0.5, 0.75, 1.0), + payroll_scale_factors=(0.5, 1.0, 1.5, 2.0), + pension_scale_factors=(0.0, 1.0), + dividend_scale_factors=(0.0, 1.0), + ), + SyntheticTemplate( + name="late_worker_single", + head_ages=(58, 60, 62, 64, 66, 68), + spouse_age_offsets=(None,), + dependent_age_sets=((),), + ss_source="older_worker", + payroll_source="prime_worker", + pension_source="prime_asset", + dividend_source="prime_asset", + ss_split=(1.0, 0.0), + payroll_split=(1.0, 0.0), + ss_scale_factors=(0.25, 0.5, 0.75, 1.0), + payroll_scale_factors=(0.5, 1.0, 1.5, 2.0), + pension_scale_factors=(0.0, 1.0), + dividend_scale_factors=(0.0, 1.0), ), ) @@ -442,7 +509,7 @@ def build_quantile_pools( "ss": _scale_levels( quantile_levels( subset["ss_total"], - quantiles=(0.25, 0.5, 0.75), + quantiles=(0.1, 0.25, 0.5, 0.75, 0.9), include_zero=(name == "zero"), positive_only=(name != "zero"), ), @@ -451,7 +518,7 @@ def build_quantile_pools( "payroll": _scale_levels( quantile_levels( subset["payroll_total"], - quantiles=(0.25, 0.5, 0.75), + quantiles=(0.1, 0.25, 0.5, 0.75, 0.9), include_zero=( name in { @@ -468,7 +535,7 @@ def build_quantile_pools( "pension": _scale_levels( quantile_levels( subset["pension_income"], - quantiles=(0.5, 0.9), + quantiles=(0.25, 0.5, 0.75, 0.9), include_zero=True, positive_only=True, ), @@ -477,7 +544,7 @@ def build_quantile_pools( "dividend": _scale_levels( quantile_levels( subset["dividend_income"], - quantiles=(0.5, 0.9), + quantiles=(0.25, 0.5, 0.75, 0.9), include_zero=True, positive_only=True, ), @@ -503,27 +570,45 @@ def generate_synthetic_candidates( ) for dependent_ages in template.dependent_age_sets: for ss_total in ss_levels: - for payroll_total in payroll_levels: - for pension_income in pension_levels: - for dividend_income in dividend_levels: - head_ss = ss_total * template.ss_split[0] - spouse_ss = ss_total * template.ss_split[1] - head_wages = payroll_total * template.payroll_split[0] - spouse_wages = payroll_total * template.payroll_split[1] - candidates.append( - SyntheticCandidate( - archetype=template.name, - head_age=head_age, - spouse_age=spouse_age, - dependent_ages=tuple(dependent_ages), - head_wages=head_wages, - spouse_wages=spouse_wages, - head_ss=head_ss, - spouse_ss=spouse_ss, - pension_income=pension_income, - dividend_income=dividend_income, - ) - ) + for ss_scale in template.ss_scale_factors: + for payroll_total in payroll_levels: + for payroll_scale in template.payroll_scale_factors: + for pension_income in pension_levels: + for pension_scale in template.pension_scale_factors: + for dividend_income in dividend_levels: + for dividend_scale in template.dividend_scale_factors: + scaled_ss_total = ss_total * ss_scale + scaled_payroll_total = payroll_total * payroll_scale + scaled_pension_income = pension_income * pension_scale + scaled_dividend_income = dividend_income * dividend_scale + head_ss = ( + scaled_ss_total * template.ss_split[0] + ) + spouse_ss = ( + scaled_ss_total * template.ss_split[1] + ) + head_wages = ( + scaled_payroll_total + * template.payroll_split[0] + ) + spouse_wages = ( + scaled_payroll_total + * template.payroll_split[1] + ) + candidates.append( + SyntheticCandidate( + archetype=template.name, + head_age=head_age, + spouse_age=spouse_age, + dependent_ages=tuple(dependent_ages), + head_wages=head_wages, + spouse_wages=spouse_wages, + head_ss=head_ss, + spouse_ss=spouse_ss, + pension_income=scaled_pension_income, + dividend_income=scaled_dividend_income, + ) + ) # Deduplicate exact duplicates caused by repeated quantiles. deduped: dict[tuple[object, ...], SyntheticCandidate] = {} for candidate in candidates: @@ -644,6 +729,24 @@ def summarize_solution( candidate_df["synthetic_weight"] / weight_sum * 100 if weight_sum > 0 else 0.0 ) candidate_df = candidate_df.sort_values("synthetic_weight", ascending=False) + positive_weights = candidate_df.loc[ + candidate_df["synthetic_weight"] > 0, + "synthetic_weight", + ].to_numpy(dtype=float) + if positive_weights.size > 0: + effective_sample_size = float( + (positive_weights.sum() ** 2) / np.sum(positive_weights**2) + ) + top_10_weight_share_pct = float( + positive_weights[:10].sum() / positive_weights.sum() * 100 + ) + top_20_weight_share_pct = float( + positive_weights[:20].sum() / positive_weights.sum() * 100 + ) + else: + effective_sample_size = 0.0 + top_10_weight_share_pct = 0.0 + top_20_weight_share_pct = 0.0 def _weighted_mean(group: pd.DataFrame, column: str) -> float: total = float(group["synthetic_weight"].sum()) @@ -757,6 +860,9 @@ def _weighted_mean(group: pd.DataFrame, column: str) -> float: return { "synthetic_candidate_count": int(len(candidate_df)), "positive_weight_candidate_count": int((candidate_df["synthetic_weight"] > 0).sum()), + "effective_sample_size": effective_sample_size, + "top_10_weight_share_pct": top_10_weight_share_pct, + "top_20_weight_share_pct": top_20_weight_share_pct, "top_candidates": candidate_df.head(20).to_dict("records"), "synthetic_archetypes": synthetic_archetypes.to_dict("records"), "actual_support_archetypes": actual_archetypes.to_dict("records"), From 2d4ab08fad2e64570236e261e922df87d5085f31 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 1 Apr 2026 08:37:00 -0400 Subject: [PATCH 21/31] Add donor-backed synthetic support probe --- .../prototype_synthetic_2100_support.py | 577 +++++++++++++++++- 1 file changed, 557 insertions(+), 20 deletions(-) diff --git a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py index 001a4418c..7c3374672 100644 --- a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py +++ b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py @@ -9,7 +9,12 @@ import pandas as pd from policyengine_us import Microsimulation -from calibration import assess_nonnegative_feasibility, calibrate_entropy +from calibration import ( + assess_nonnegative_feasibility, + calibrate_entropy, + calibrate_entropy_bounded, + densify_lp_solution, +) from projection_utils import aggregate_age_targets, build_age_bins from ssa_data import ( get_long_term_target_source, @@ -306,9 +311,35 @@ def parse_args() -> argparse.Namespace: type=Path, help="Optional JSON output path.", ) + parser.add_argument( + "--epsilon-path", + default="0.25,0.5,1.0", + help=( + "Comma-separated approximate error thresholds to probe with " + "bounded entropy after the exact solve. Use an empty string to disable." + ), + ) + parser.add_argument( + "--donor-probe-top-n", + type=int, + default=20, + help="Number of top exact-fit synthetic candidates to probe against real donors.", + ) + parser.add_argument( + "--donor-probe-k", + type=int, + default=5, + help="Number of nearest real donors to report for each probed synthetic candidate.", + ) return parser.parse_args() +def parse_epsilon_path(value: str) -> tuple[float, ...]: + if not value.strip(): + return () + return tuple(float(part.strip()) for part in value.split(",") if part.strip()) + + def classify_archetype( *, head_age: float, @@ -394,11 +425,23 @@ def build_actual_tax_unit_summary(base_dataset: str) -> pd.DataFrame: float(spouses["age"].iloc[0]) if not spouses.empty else None ) dependent_count = int(len(dependents)) + dependent_ages = tuple(sorted(int(age) for age in dependents["age"].tolist())) + head_payroll = float(heads["payroll"].sum()) if not heads.empty else 0.0 + spouse_payroll = float(spouses["payroll"].sum()) if not spouses.empty else 0.0 + head_ss = float(heads["social_security"].sum()) if not heads.empty else 0.0 + spouse_ss = ( + float(spouses["social_security"].sum()) if not spouses.empty else 0.0 + ) row = { "tax_unit_id": int(tax_unit_id), "head_age": head_age, "spouse_age": spouse_age, "dependent_count": dependent_count, + "dependent_ages": dependent_ages, + "head_payroll": head_payroll, + "spouse_payroll": spouse_payroll, + "head_ss": head_ss, + "spouse_ss": spouse_ss, "payroll_total": float(group["payroll"].sum()), "ss_total": float(group["social_security"].sum()), "dividend_income": float(group["dividend_income"].sum()), @@ -641,21 +684,428 @@ def age_bucket_vector(ages: list[int], age_bins: list[tuple[int, int]]) -> np.nd return vector -def solve_synthetic_support( +def build_synthetic_constraint_problem( candidates: list[SyntheticCandidate], *, year: int, -) -> tuple[np.ndarray, dict[str, object]]: +) -> dict[str, object]: age_targets = load_ssa_age_projections(start_year=year, end_year=year) age_bins = build_age_bins(n_ages=age_targets.shape[0], bucket_size=5) aggregated_age_targets = aggregate_age_targets(age_targets, age_bins)[:, 0] - X = np.vstack([age_bucket_vector(candidate.ages(), age_bins) for candidate in candidates]) ss_values = np.array([candidate.ss_total for candidate in candidates], dtype=float) payroll_values = np.array( - [candidate.payroll_total for candidate in candidates], dtype=float + [candidate.payroll_total for candidate in candidates], + dtype=float, ) - baseline_weights = np.ones(len(candidates), dtype=float) + return { + "age_bins": age_bins, + "aggregated_age_targets": aggregated_age_targets, + "X": X, + "ss_values": ss_values, + "ss_target": float(load_ssa_benefit_projections(year)), + "payroll_values": payroll_values, + "payroll_target": float(load_taxable_payroll_projections(year)), + "baseline_weights": np.ones(len(candidates), dtype=float), + } + + +def build_constraint_matrix(problem: dict[str, object]) -> tuple[np.ndarray, np.ndarray]: + constraint_matrix = np.column_stack( + [ + problem["X"], + problem["ss_values"], + problem["payroll_values"], + ] + ) + targets = np.concatenate( + [ + problem["aggregated_age_targets"], + np.array([problem["ss_target"], problem["payroll_target"]], dtype=float), + ] + ) + return constraint_matrix, targets + + +def build_scaled_actual_summary( + actual_summary: pd.DataFrame, + *, + ss_scale: float, + earnings_scale: float, +) -> pd.DataFrame: + scaled = actual_summary.copy() + scaled["scaled_head_payroll"] = scaled["head_payroll"] * earnings_scale + scaled["scaled_spouse_payroll"] = scaled["spouse_payroll"] * earnings_scale + scaled["scaled_payroll_total"] = scaled["payroll_total"] * earnings_scale + scaled["scaled_head_ss"] = scaled["head_ss"] * ss_scale + scaled["scaled_spouse_ss"] = scaled["spouse_ss"] * ss_scale + scaled["scaled_ss_total"] = scaled["ss_total"] * ss_scale + scaled["spouse_present"] = scaled["spouse_age"].notna() + scaled["spouse_age_filled"] = scaled["spouse_age"].fillna(-1) + return scaled + + +def summarize_exact_candidates( + candidates: list[SyntheticCandidate], + weights: np.ndarray, +) -> pd.DataFrame: + candidate_rows = [] + for idx, (candidate, weight) in enumerate(zip(candidates, weights)): + candidate_rows.append( + { + "candidate_idx": idx, + **asdict(candidate), + "dependent_count": len(candidate.dependent_ages), + "payroll_total": candidate.payroll_total, + "ss_total": candidate.ss_total, + "synthetic_weight": float(weight), + } + ) + candidate_df = pd.DataFrame(candidate_rows).sort_values( + "synthetic_weight", + ascending=False, + ) + total_weight = max(float(candidate_df["synthetic_weight"].sum()), 1.0) + candidate_df["weight_share_pct"] = ( + candidate_df["synthetic_weight"] / total_weight * 100 + ) + return candidate_df + + +def match_real_donors_for_target( + target_row: pd.Series, + scaled_actual_summary: pd.DataFrame, + *, + donors_per_target: int, +) -> pd.DataFrame: + target_spouse_present = pd.notna(target_row["spouse_age"]) + target_spouse_age = -1 if not target_spouse_present else target_row["spouse_age"] + subset = scaled_actual_summary[ + scaled_actual_summary["spouse_present"].eq(target_spouse_present) + ].copy() + if subset.empty: + subset = scaled_actual_summary.copy() + subset["distance"] = ( + (subset["head_age"] - target_row["head_age"]).abs() / 5.0 + + ( + subset["spouse_age_filled"] - target_spouse_age + ).abs() + / 5.0 + + (subset["dependent_count"] - target_row["dependent_count"]).abs() * 0.75 + + np.abs( + np.log1p(subset["scaled_payroll_total"]) + - np.log1p(float(target_row["payroll_total"])) + ) + + np.abs( + np.log1p(subset["scaled_ss_total"]) + - np.log1p(float(target_row["ss_total"])) + ) + + 0.25 * subset["archetype"].ne(target_row["archetype"]).astype(float) + ) + nearest = subset.nsmallest(donors_per_target, "distance").copy() + nearest["target_candidate_idx"] = int(target_row["candidate_idx"]) + nearest["target_archetype"] = target_row["archetype"] + nearest["target_weight_share_pct"] = float(target_row["weight_share_pct"]) + nearest["target_head_age"] = int(target_row["head_age"]) + nearest["target_spouse_age"] = ( + None if not target_spouse_present else int(target_row["spouse_age"]) + ) + nearest["target_dependent_count"] = int(target_row["dependent_count"]) + nearest["target_payroll_total"] = float(target_row["payroll_total"]) + nearest["target_ss_total"] = float(target_row["ss_total"]) + nearest["required_head_age_shift"] = ( + nearest["target_head_age"] - nearest["head_age"] + ) + nearest["required_spouse_age_shift"] = np.where( + nearest["target_spouse_age"].isna(), + np.nan, + nearest["target_spouse_age"] - nearest["spouse_age_filled"], + ) + return nearest + + +def summarize_donor_probe( + candidates: list[SyntheticCandidate], + weights: np.ndarray, + actual_summary: pd.DataFrame, + *, + ss_scale: float, + earnings_scale: float, + top_n_targets: int, + donors_per_target: int, +) -> dict[str, object]: + exact_df = summarize_exact_candidates(candidates, weights) + target_df = exact_df[exact_df["synthetic_weight"] > 0].head(top_n_targets).copy() + scaled_actual = build_scaled_actual_summary( + actual_summary, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + ) + donor_matches = [] + for _, target_row in target_df.iterrows(): + donor_matches.append( + match_real_donors_for_target( + target_row, + scaled_actual, + donors_per_target=donors_per_target, + ) + ) + donor_df = pd.concat(donor_matches, ignore_index=True) + nearest_only = donor_df.sort_values( + ["target_candidate_idx", "distance"], + ascending=[True, True], + ).groupby("target_candidate_idx", as_index=False).first() + distance_summary = { + "median_best_distance": float(nearest_only["distance"].median()), + "targets_with_best_distance_le_1": int((nearest_only["distance"] <= 1.0).sum()), + "targets_with_best_distance_le_2": int((nearest_only["distance"] <= 2.0).sum()), + "targets_with_best_distance_gt_3": int((nearest_only["distance"] > 3.0).sum()), + } + outlier_targets = nearest_only[nearest_only["distance"] > 3.0].copy() + return { + "top_n_targets": int(top_n_targets), + "donors_per_target": int(donors_per_target), + "distance_summary": distance_summary, + "nearest_targets": nearest_only[ + [ + "target_candidate_idx", + "target_archetype", + "target_weight_share_pct", + "target_head_age", + "target_spouse_age", + "target_dependent_count", + "target_payroll_total", + "target_ss_total", + "tax_unit_id", + "archetype", + "head_age", + "spouse_age", + "dependent_count", + "scaled_payroll_total", + "scaled_ss_total", + "required_head_age_shift", + "required_spouse_age_shift", + "distance", + ] + ].to_dict("records"), + "outlier_targets": outlier_targets[ + [ + "target_candidate_idx", + "target_archetype", + "target_weight_share_pct", + "target_head_age", + "target_spouse_age", + "target_dependent_count", + "target_payroll_total", + "target_ss_total", + "distance", + ] + ].to_dict("records"), + } + + +def _safe_split(numerator: float, denominator: float, fallback: float) -> float: + if denominator <= 0: + return fallback + return float(numerator / denominator) + + +def build_donor_backed_clones( + candidates: list[SyntheticCandidate], + weights: np.ndarray, + actual_summary: pd.DataFrame, + *, + ss_scale: float, + earnings_scale: float, + top_n_targets: int, + donors_per_target: int, + max_distance_for_clone: float = 3.0, +) -> tuple[list[SyntheticCandidate], np.ndarray, dict[str, object]]: + exact_df = summarize_exact_candidates(candidates, weights) + target_df = exact_df[exact_df["synthetic_weight"] > 0].head(top_n_targets).copy() + scaled_actual = build_scaled_actual_summary( + actual_summary, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + ) + + donor_backed_candidates: list[SyntheticCandidate] = [] + donor_backed_weights: list[float] = [] + clone_records: list[dict[str, object]] = [] + outlier_targets: list[dict[str, object]] = [] + + for _, target_row in target_df.iterrows(): + target_candidate = candidates[int(target_row["candidate_idx"])] + donor_matches = match_real_donors_for_target( + target_row, + scaled_actual, + donors_per_target=donors_per_target, + ) + usable = donor_matches[donor_matches["distance"] <= max_distance_for_clone].copy() + if usable.empty: + donor_backed_candidates.append(target_candidate) + donor_backed_weights.append(float(target_row["synthetic_weight"])) + outlier_targets.append( + { + "candidate_idx": int(target_row["candidate_idx"]), + "archetype": target_row["archetype"], + "weight_share_pct": float(target_row["weight_share_pct"]), + "best_distance": float(donor_matches["distance"].min()), + } + ) + continue + + per_clone_weight = float(target_row["synthetic_weight"]) / len(usable) + target_payroll_total = float(target_row["payroll_total"]) + target_ss_total = float(target_row["ss_total"]) + target_head_payroll_share = _safe_split( + target_candidate.head_wages, + target_candidate.payroll_total, + 1.0 if target_candidate.spouse_age is None else 0.5, + ) + target_head_ss_share = _safe_split( + target_candidate.head_ss, + target_candidate.ss_total, + 1.0 if target_candidate.spouse_age is None else 0.5, + ) + + for _, donor_row in usable.iterrows(): + donor_head_payroll_share = _safe_split( + float(donor_row["scaled_head_payroll"]), + float(donor_row["scaled_payroll_total"]), + target_head_payroll_share, + ) + donor_head_ss_share = _safe_split( + float(donor_row["scaled_head_ss"]), + float(donor_row["scaled_ss_total"]), + target_head_ss_share, + ) + donor_backed_candidates.append( + SyntheticCandidate( + archetype=target_candidate.archetype, + head_age=target_candidate.head_age, + spouse_age=target_candidate.spouse_age, + dependent_ages=target_candidate.dependent_ages, + head_wages=target_payroll_total * donor_head_payroll_share, + spouse_wages=target_payroll_total * (1.0 - donor_head_payroll_share), + head_ss=target_ss_total * donor_head_ss_share, + spouse_ss=target_ss_total * (1.0 - donor_head_ss_share), + pension_income=float(donor_row["pension_income"]) * earnings_scale, + dividend_income=float(donor_row["dividend_income"]) * earnings_scale, + ) + ) + donor_backed_weights.append(per_clone_weight) + clone_records.append( + { + "candidate_idx": int(target_row["candidate_idx"]), + "archetype": target_row["archetype"], + "tax_unit_id": int(donor_row["tax_unit_id"]), + "distance": float(donor_row["distance"]), + "assigned_weight_share_pct": float( + per_clone_weight / max(float(weights.sum()), 1.0) * 100 + ), + } + ) + + clone_summary = summarize_solution( + donor_backed_candidates, + np.asarray(donor_backed_weights, dtype=float), + actual_summary, + ) + return ( + donor_backed_candidates, + np.asarray(donor_backed_weights, dtype=float), + { + "top_n_targets": int(top_n_targets), + "donors_per_target": int(donors_per_target), + "max_distance_for_clone": float(max_distance_for_clone), + "outlier_targets": outlier_targets, + "clone_records": clone_records[:100], + "clone_summary": clone_summary, + }, + ) + + +def solve_synthetic_support( + candidates: list[SyntheticCandidate], + *, + year: int, + max_constraint_error_pct: float = 0.0, + warm_weights: np.ndarray | None = None, +) -> tuple[np.ndarray, dict[str, object]]: + problem = build_synthetic_constraint_problem(candidates, year=year) + X = problem["X"] + aggregated_age_targets = problem["aggregated_age_targets"] + ss_values = problem["ss_values"] + payroll_values = problem["payroll_values"] + baseline_weights = problem["baseline_weights"] + ss_target = problem["ss_target"] + payroll_target = problem["payroll_target"] + + if max_constraint_error_pct > 0: + try: + weights, iterations, info = calibrate_entropy_bounded( + X, + aggregated_age_targets, + baseline_weights, + ss_values=ss_values, + ss_target=ss_target, + payroll_values=payroll_values, + payroll_target=payroll_target, + n_ages=X.shape[1], + max_constraint_error_pct=max_constraint_error_pct, + max_iters=500, + tol=1e-9, + warm_weights=([warm_weights] if warm_weights is not None else None), + ) + return np.asarray(weights, dtype=float), { + "method": "bounded_entropy", + "iterations": int(iterations), + "best_case_max_pct_error": float(info["best_case_max_pct_error"]), + "requested_max_constraint_error_pct": float(max_constraint_error_pct), + "age_bucket_size": 5, + "status": int(info.get("status", 0)), + "message": info.get("message"), + } + except RuntimeError as error: + constraint_matrix, targets = build_constraint_matrix(problem) + lp_weights = None + if warm_weights is not None: + lp_weights = np.asarray(warm_weights, dtype=float) + else: + feasibility = assess_nonnegative_feasibility( + constraint_matrix, + targets, + return_weights=True, + ) + if feasibility["success"] and feasibility.get("weights") is not None: + lp_weights = np.asarray(feasibility["weights"], dtype=float) + if lp_weights is None: + raise RuntimeError( + f"Approximate synthetic support solve failed for {year}: {error}" + ) from error + dense_weights, dense_info = densify_lp_solution( + constraint_matrix, + targets, + baseline_weights, + lp_weights, + max_constraint_error_pct, + ) + return np.asarray(dense_weights, dtype=float), { + "method": ( + "lp_blend" + if dense_info["densification_effective"] + else "lp_minimax" + ), + "iterations": 1, + "best_case_max_pct_error": float( + dense_info["best_case_max_pct_error"] + ), + "requested_max_constraint_error_pct": float(max_constraint_error_pct), + "age_bucket_size": 5, + "entropy_error": str(error), + "lp_blend_lambda": float(dense_info["blend_lambda"]), + } try: weights, iterations = calibrate_entropy( @@ -663,9 +1113,9 @@ def solve_synthetic_support( aggregated_age_targets, baseline_weights, ss_values=ss_values, - ss_target=load_ssa_benefit_projections(year), + ss_target=ss_target, payroll_values=payroll_values, - payroll_target=load_taxable_payroll_projections(year), + payroll_target=payroll_target, n_ages=X.shape[1], max_iters=500, tol=1e-9, @@ -677,19 +1127,10 @@ def solve_synthetic_support( "age_bucket_size": 5, } except RuntimeError as error: + constraint_matrix, targets = build_constraint_matrix(problem) feasibility = assess_nonnegative_feasibility( - np.column_stack([X, ss_values, payroll_values]), - np.concatenate( - [ - aggregated_age_targets, - np.array( - [ - load_ssa_benefit_projections(year), - load_taxable_payroll_projections(year), - ] - ), - ] - ), + constraint_matrix, + targets, return_weights=True, ) if not feasibility["success"] or feasibility.get("weights") is None: @@ -705,6 +1146,57 @@ def solve_synthetic_support( } +def summarize_solution_diff( + candidates: list[SyntheticCandidate], + base_weights: np.ndarray, + alt_weights: np.ndarray, +) -> dict[str, object]: + base_weights = np.asarray(base_weights, dtype=float) + alt_weights = np.asarray(alt_weights, dtype=float) + base_total = max(float(base_weights.sum()), 1.0) + alt_total = max(float(alt_weights.sum()), 1.0) + rows: list[dict[str, object]] = [] + for candidate, base_weight, alt_weight in zip(candidates, base_weights, alt_weights): + base_share_pct = float(base_weight / base_total * 100) + alt_share_pct = float(alt_weight / alt_total * 100) + rows.append( + { + "archetype": candidate.archetype, + "head_age": candidate.head_age, + "spouse_age": candidate.spouse_age, + "dependent_count": len(candidate.dependent_ages), + "payroll_total": candidate.payroll_total, + "ss_total": candidate.ss_total, + "base_weight_share_pct": base_share_pct, + "alt_weight_share_pct": alt_share_pct, + "weight_share_gain_pct_points": alt_share_pct - base_share_pct, + "newly_entering": alt_weight > 0 and base_weight <= 1e-12, + } + ) + diff_df = pd.DataFrame(rows) + entrants = diff_df[ + diff_df["newly_entering"] & diff_df["alt_weight_share_pct"].gt(0.01) + ].copy() + entrant_archetypes = ( + entrants.groupby("archetype", as_index=False) + .agg( + entrant_weight_share_pct=("alt_weight_share_pct", "sum"), + entrant_candidate_count=("archetype", "count"), + ) + .sort_values("entrant_weight_share_pct", ascending=False) + ) + return { + "top_weight_gainers": diff_df.sort_values( + "weight_share_gain_pct_points", + ascending=False, + ) + .head(20) + .to_dict("records"), + "entrant_archetypes": entrant_archetypes.head(12).to_dict("records"), + "entrant_candidate_count": int(len(entrants)), + } + + def summarize_solution( candidates: list[SyntheticCandidate], weights: np.ndarray, @@ -893,6 +1385,48 @@ def main() -> int: candidates = generate_synthetic_candidates(pools) weights, solve_info = solve_synthetic_support(candidates, year=args.year) solution_summary = summarize_solution(candidates, weights, actual_summary) + donor_probe = summarize_donor_probe( + candidates, + weights, + actual_summary, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + top_n_targets=args.donor_probe_top_n, + donors_per_target=args.donor_probe_k, + ) + _, _, donor_backed_clone_probe = build_donor_backed_clones( + candidates, + weights, + actual_summary, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + top_n_targets=args.donor_probe_top_n, + donors_per_target=args.donor_probe_k, + ) + epsilon_path_results = [] + for epsilon in parse_epsilon_path(args.epsilon_path): + epsilon_weights, epsilon_solve_info = solve_synthetic_support( + candidates, + year=args.year, + max_constraint_error_pct=epsilon, + warm_weights=weights, + ) + epsilon_path_results.append( + { + "epsilon_pct": float(epsilon), + "solve_info": epsilon_solve_info, + "solution_summary": summarize_solution( + candidates, + epsilon_weights, + actual_summary, + ), + "vs_exact": summarize_solution_diff( + candidates, + weights, + epsilon_weights, + ), + } + ) report = { "year": args.year, @@ -910,6 +1444,9 @@ def main() -> int: }, "actual_support_tax_unit_count": int(len(actual_summary)), "synthetic_solution": solution_summary, + "donor_probe": donor_probe, + "donor_backed_clone_probe": donor_backed_clone_probe, + "epsilon_path": epsilon_path_results, } payload = json.dumps(report, indent=2) From de6403db48df5c0545ecaae60cd3533221836cdc Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 1 Apr 2026 09:17:46 -0400 Subject: [PATCH 22/31] Add donor-backed late-year support mode --- .../datasets/cps/long_term/README.md | 21 ++ .../cps/long_term/calibration_artifacts.py | 17 + .../cps/long_term/projection_utils.py | 6 +- .../prototype_synthetic_2100_support.py | 334 ++++++++++++++++++ .../cps/long_term/run_household_projection.py | 176 ++++++++- .../test_long_term_calibration_contract.py | 58 +++ 6 files changed, 605 insertions(+), 7 deletions(-) diff --git a/policyengine_us_data/datasets/cps/long_term/README.md b/policyengine_us_data/datasets/cps/long_term/README.md index 11f082824..dc361120e 100644 --- a/policyengine_us_data/datasets/cps/long_term/README.md +++ b/policyengine_us_data/datasets/cps/long_term/README.md @@ -9,6 +9,9 @@ Run projections using `run_household_projection.py`: # Recommended: named profile with TOB benchmarked post-calibration python run_household_projection.py 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --save-h5 +# Experimental: donor-backed late-year support augmentation for tail-year runs +python run_household_projection.py 2075 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --support-augmentation-profile donor-backed-synthetic-v1 --support-augmentation-target-year 2100 --allow-validation-failures + # IPF with only age distribution constraints (faster, less accurate) python run_household_projection.py 2050 --profile age-only @@ -21,6 +24,13 @@ python run_household_projection.py 2100 --profile ss - `--profile`: Named calibration contract. Recommended over legacy flags. - `--target-source`: Named long-term target source package. - `--output-dir`: Output directory for generated H5 files and metadata sidecars. +- `--support-augmentation-profile`: Experimental late-year support expansion mode. Currently supports `donor-backed-synthetic-v1`. +- `--support-augmentation-target-year`: Extreme year used to build the donor-backed supplement (defaults to `END_YEAR`). +- `--support-augmentation-start-year`: Earliest run year allowed for augmentation (defaults to `2075`). +- `--support-augmentation-top-n-targets`: Number of dominant synthetic target types to map back to real donors (default `20`). +- `--support-augmentation-donors-per-target`: Number of nearest real donor tax units per synthetic target (default `5`). +- `--support-augmentation-max-distance`: Maximum donor-match distance retained for cloning (default `3.0`). +- `--support-augmentation-clone-weight-scale`: Baseline weight multiplier applied to each donor-backed clone (default `0.1`). - `--greg`: Use GREG calibration instead of IPF - `--use-ss`: Include Social Security benefit totals as calibration target (requires `--greg`) - `--use-payroll`: Include taxable payroll totals as calibration target (requires `--greg`) @@ -42,6 +52,7 @@ python run_household_projection.py 2100 --profile ss - Each output directory now gets a `calibration_manifest.json` file describing the profile/base dataset contract for the full artifact set. - Profiles validate achieved constraint errors before writing output. +- Experimental donor-backed augmentation is stamped into each year sidecar and the directory manifest via `support_augmentation`. **Estimated runtime:** ~2 minutes/year without `--save-h5`, ~3 minutes/year with `--save-h5` @@ -60,6 +71,14 @@ python run_household_projection.py 2100 --profile ss - Uses dual optimization to minimize divergence from baseline weights - **Recommended** for publishable long-term projections +**Donor-Backed Late-Year Support Augmentation** +- Experimental late-tail option for `2075+` runs +- Uses the `2100` synthetic-support prototype to identify dominant missing household types +- Maps those synthetic targets back to nearest real 2024 donor tax units +- Clones and perturbs the donor tax units to create a small augmented support without replacing the base CPS sample +- Intended to test whether donor-backed synthetic support improves late-year microsim feasibility without resorting to fully free synthetic records +- Current status: integrated into the runner and fully auditable in metadata, but still diagnostic. The first `2100` end-to-end run did not materially improve the late-tail calibration frontier or support-concentration metrics. + **GREG (Generalized Regression Estimator)** - Legacy linear calibration path retained for explicit flag-based runs - Can hit constraints exactly, but may produce negative weights in far-horizon years @@ -146,6 +165,8 @@ For each projection year (2025-2100): 5. **Benchmark TOB** - Compare modeled OASDI/HI TOB to the selected target source without forcing it into the weights 6. **Aggregate results** - Apply calibrated weights to calculate national totals +When donor-backed augmentation is enabled, step 1 uses the original 2024 CPS support and step 2 inserts a tagged late-year supplement derived from nearest real donors before the calibration loop begins. The underlying base dataset path remains unchanged in metadata; the augmentation details are recorded separately in `support_augmentation`. + **Key innovation:** Household-level calculations avoid person→household aggregation issues, maintaining consistency across all variables. --- diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py index 5be53d18d..2676895b9 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py @@ -102,6 +102,7 @@ def write_year_metadata( profile: dict[str, Any], calibration_audit: dict[str, Any], target_source: dict[str, Any] | None = None, + support_augmentation: dict[str, Any] | None = None, ) -> Path: metadata = { "contract_version": CONTRACT_VERSION, @@ -112,6 +113,8 @@ def write_year_metadata( } if target_source is not None: metadata["target_source"] = target_source + if support_augmentation is not None: + metadata["support_augmentation"] = support_augmentation metadata = normalize_metadata(metadata) metadata_path = metadata_path_for(h5_path) metadata_path.write_text( @@ -131,11 +134,13 @@ def update_dataset_manifest( profile: dict[str, Any], calibration_audit: dict[str, Any], target_source: dict[str, Any] | None = None, + support_augmentation: dict[str, Any] | None = None, ) -> Path: output_dir = Path(output_dir) manifest_path = output_dir / MANIFEST_FILENAME profile = json.loads(json.dumps(profile)) target_source = json.loads(json.dumps(target_source)) + support_augmentation = json.loads(json.dumps(support_augmentation)) if manifest_path.exists(): manifest = json.loads(manifest_path.read_text(encoding="utf-8")) @@ -146,6 +151,7 @@ def update_dataset_manifest( "base_dataset_path": base_dataset_path, "profile": profile, "target_source": target_source, + "support_augmentation": support_augmentation, "years": [], "datasets": {}, } @@ -175,6 +181,16 @@ def update_dataset_manifest( "Output directory already contains a different target source: " f"{manifest.get('target_source')} != {target_source}" ) + if ( + manifest.get("support_augmentation") is None + and support_augmentation is not None + ): + manifest["support_augmentation"] = support_augmentation + elif manifest.get("support_augmentation") != support_augmentation: + raise ValueError( + "Output directory already contains a different support augmentation: " + f"{manifest.get('support_augmentation')} != {support_augmentation}" + ) datasets = manifest.setdefault("datasets", {}) datasets[str(year)] = { @@ -251,6 +267,7 @@ def rebuild_dataset_manifest_with_target_source( profile=metadata["profile"], calibration_audit=metadata["calibration_audit"], target_source=metadata.get("target_source"), + support_augmentation=metadata.get("support_augmentation"), ) assert manifest_path is not None diff --git a/policyengine_us_data/datasets/cps/long_term/projection_utils.py b/policyengine_us_data/datasets/cps/long_term/projection_utils.py index 390dcc690..7c75122f6 100644 --- a/policyengine_us_data/datasets/cps/long_term/projection_utils.py +++ b/policyengine_us_data/datasets/cps/long_term/projection_utils.py @@ -115,7 +115,7 @@ def get_pseudo_input_variables(sim): return pseudo_inputs -def create_household_year_h5(year, household_weights, base_dataset_path, output_dir): +def create_household_year_h5(year, household_weights, base_dataset, output_dir): """ Create a year-specific .h5 file with calibrated household weights. @@ -125,7 +125,7 @@ def create_household_year_h5(year, household_weights, base_dataset_path, output_ Args: year: The year for this dataset household_weights: Calibrated household weights for this year - base_dataset_path: Path to base dataset + base_dataset: Path to base dataset or in-memory Dataset instance output_dir: Directory to save the .h5 file Returns: @@ -133,7 +133,7 @@ def create_household_year_h5(year, household_weights, base_dataset_path, output_ """ output_path = os.path.join(output_dir, f"{year}.h5") - sim = Microsimulation(dataset=base_dataset_path) + sim = Microsimulation(dataset=base_dataset) base_period = int(sim.default_calculation_period) df = sim.to_input_dataframe() diff --git a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py index 7c3374672..12bcec437 100644 --- a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py +++ b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py @@ -7,6 +7,7 @@ import numpy as np import pandas as pd +from policyengine_core.data.dataset import Dataset from policyengine_us import Microsimulation from calibration import ( @@ -28,6 +29,24 @@ DEFAULT_DATASET = "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5" DEFAULT_YEAR = 2100 BASE_YEAR = 2024 +ENTITY_ID_COLUMNS = { + "household": ("household_id", "person_household_id"), + "family": ("family_id", "person_family_id"), + "tax_unit": ("tax_unit_id", "person_tax_unit_id"), + "spm_unit": ("spm_unit_id", "person_spm_unit_id"), + "marital_unit": ("marital_unit_id", "person_marital_unit_id"), +} +PERSON_ID_COLUMN = "person_id" +SS_COMPONENTS = ( + "social_security_retirement", + "social_security_disability", + "social_security_survivors", + "social_security_dependents", +) +PAYROLL_COMPONENTS = ( + "employment_income_before_lsr", + "self_employment_income_before_lsr", +) @dataclass(frozen=True) @@ -340,6 +359,10 @@ def parse_epsilon_path(value: str) -> tuple[float, ...]: return tuple(float(part.strip()) for part in value.split(",") if part.strip()) +def _period_column(name: str, base_year: int) -> str: + return f"{name}__{base_year}" + + def classify_archetype( *, head_age: float, @@ -425,6 +448,7 @@ def build_actual_tax_unit_summary(base_dataset: str) -> pd.DataFrame: float(spouses["age"].iloc[0]) if not spouses.empty else None ) dependent_count = int(len(dependents)) + adult_count = int((group["age"] >= 18).sum()) dependent_ages = tuple(sorted(int(age) for age in dependents["age"].tolist())) head_payroll = float(heads["payroll"].sum()) if not heads.empty else 0.0 spouse_payroll = float(spouses["payroll"].sum()) if not spouses.empty else 0.0 @@ -436,6 +460,7 @@ def build_actual_tax_unit_summary(base_dataset: str) -> pd.DataFrame: "tax_unit_id": int(tax_unit_id), "head_age": head_age, "spouse_age": spouse_age, + "adult_count": adult_count, "dependent_count": dependent_count, "dependent_ages": dependent_ages, "head_payroll": head_payroll, @@ -780,9 +805,16 @@ def match_real_donors_for_target( ) -> pd.DataFrame: target_spouse_present = pd.notna(target_row["spouse_age"]) target_spouse_age = -1 if not target_spouse_present else target_row["spouse_age"] + target_adult_count = 2 if target_spouse_present else 1 subset = scaled_actual_summary[ scaled_actual_summary["spouse_present"].eq(target_spouse_present) + & scaled_actual_summary["adult_count"].eq(target_adult_count) + & scaled_actual_summary["dependent_count"].eq(int(target_row["dependent_count"])) ].copy() + if subset.empty: + subset = scaled_actual_summary[ + scaled_actual_summary["adult_count"].ge(target_adult_count) + ].copy() if subset.empty: subset = scaled_actual_summary.copy() subset["distance"] = ( @@ -904,6 +936,308 @@ def summarize_donor_probe( } +def _next_entity_id(values: pd.Series) -> int: + non_null = values.dropna() + if non_null.empty: + return 1 + return int(non_null.max()) + 1 + + +def _cast_mapped_ids(series: pd.Series, mapped: pd.Series) -> pd.Series: + dtype = series.dtype + if pd.api.types.is_integer_dtype(dtype): + return mapped.astype(dtype) + if pd.api.types.is_float_dtype(dtype): + return mapped.astype(dtype) + return mapped + + +def _scale_person_components( + row: pd.Series, + columns: tuple[str, ...], + target_total: float, +) -> pd.Series: + available = [column for column in columns if column in row.index] + if not available: + return row + target_total = float(target_total) + if target_total <= 0: + for column in available: + row[column] = 0.0 + return row + current_total = float(sum(float(row[column]) for column in available)) + if current_total > 0: + scale = target_total / current_total + for column in available: + row[column] = float(row[column]) * scale + return row + row[available[0]] = target_total + for column in available[1:]: + row[column] = 0.0 + return row + + +def _clone_tax_unit_rows_to_target( + donor_rows: pd.DataFrame, + *, + base_year: int, + target_candidate: SyntheticCandidate, + ss_scale: float, + earnings_scale: float, + id_counters: dict[str, int], + clone_weight_scale: float, + clone_weight_divisor: int, +) -> tuple[pd.DataFrame, dict[str, int]] | tuple[None, dict[str, int]]: + age_col = _period_column("age", base_year) + household_weight_col = _period_column("household_weight", base_year) + person_weight_col = _period_column("person_weight", base_year) + person_id_col = _period_column(PERSON_ID_COLUMN, base_year) + + adults = donor_rows[donor_rows[age_col] >= 18].sort_values(age_col, ascending=False) + dependents = donor_rows[donor_rows[age_col] < 18].sort_values(age_col, ascending=False) + target_has_spouse = target_candidate.spouse_age is not None + target_adult_count = 2 if target_has_spouse else 1 + if len(adults) < target_adult_count or len(dependents) != len(target_candidate.dependent_ages): + return None, id_counters + + cloned = donor_rows.copy() + household_id = id_counters["household"] + id_counters["household"] += 1 + for entity_name, columns in ENTITY_ID_COLUMNS.items(): + entity_id = id_counters[entity_name] + id_counters[entity_name] += 1 + for raw_column in columns: + column = _period_column(raw_column, base_year) + if column in cloned.columns: + cloned[column] = entity_id if entity_name != "household" else household_id + cloned[_period_column("household_id", base_year)] = household_id + cloned[_period_column("person_household_id", base_year)] = household_id + + person_ids = range(id_counters["person"], id_counters["person"] + len(cloned)) + id_counters["person"] += len(cloned) + cloned[person_id_col] = _cast_mapped_ids( + cloned[person_id_col], + pd.Series(list(person_ids), index=cloned.index), + ) + + if household_weight_col in cloned.columns: + cloned[household_weight_col] = ( + cloned[household_weight_col].astype(float) + * clone_weight_scale + / max(clone_weight_divisor, 1) + ) + if person_weight_col in cloned.columns: + cloned[person_weight_col] = ( + cloned[person_weight_col].astype(float) + * clone_weight_scale + / max(clone_weight_divisor, 1) + ) + + adult_indices = adults.index.tolist() + head_idx = adult_indices[0] + spouse_idx = adult_indices[1] if target_has_spouse else None + dependent_indices = dependents.index.tolist() + + cloned.loc[head_idx, age_col] = float(target_candidate.head_age) + if spouse_idx is not None: + cloned.loc[spouse_idx, age_col] = float(target_candidate.spouse_age) + for dep_idx, dep_age in zip(dependent_indices, target_candidate.dependent_ages): + cloned.loc[dep_idx, age_col] = float(dep_age) + + payroll_columns = tuple(_period_column(component, base_year) for component in PAYROLL_COMPONENTS) + ss_columns = tuple(_period_column(component, base_year) for component in SS_COMPONENTS) + qbi_col = _period_column("w2_wages_from_qualified_business", base_year) + + target_head_payroll = float(target_candidate.head_wages) / max(earnings_scale, 1e-12) + target_spouse_payroll = float(target_candidate.spouse_wages) / max(earnings_scale, 1e-12) + target_head_ss = float(target_candidate.head_ss) / max(ss_scale, 1e-12) + target_spouse_ss = float(target_candidate.spouse_ss) / max(ss_scale, 1e-12) + + cloned.loc[head_idx] = _scale_person_components( + cloned.loc[head_idx].copy(), + payroll_columns, + target_head_payroll, + ) + cloned.loc[head_idx] = _scale_person_components( + cloned.loc[head_idx].copy(), + ss_columns, + target_head_ss, + ) + if spouse_idx is not None: + cloned.loc[spouse_idx] = _scale_person_components( + cloned.loc[spouse_idx].copy(), + payroll_columns, + target_spouse_payroll, + ) + cloned.loc[spouse_idx] = _scale_person_components( + cloned.loc[spouse_idx].copy(), + ss_columns, + target_spouse_ss, + ) + + for dep_idx in dependent_indices: + cloned.loc[dep_idx] = _scale_person_components( + cloned.loc[dep_idx].copy(), + payroll_columns, + 0.0, + ) + cloned.loc[dep_idx] = _scale_person_components( + cloned.loc[dep_idx].copy(), + ss_columns, + 0.0, + ) + if qbi_col in cloned.columns: + cloned.loc[dep_idx, qbi_col] = 0.0 + + if qbi_col in cloned.columns and head_idx in cloned.index: + cloned.loc[head_idx, qbi_col] = 0.0 + if spouse_idx is not None: + cloned.loc[spouse_idx, qbi_col] = 0.0 + + return cloned, id_counters + + +def build_donor_backed_augmented_input_dataframe( + *, + base_dataset: str, + base_year: int, + target_year: int, + top_n_targets: int = 20, + donors_per_target: int = 5, + max_distance_for_clone: float = 3.0, + clone_weight_scale: float = 0.1, +) -> tuple[pd.DataFrame, dict[str, object]]: + sim = Microsimulation(dataset=base_dataset) + input_df = sim.to_input_dataframe() + actual_summary = build_actual_tax_unit_summary(base_dataset) + base_aggregates = load_base_aggregates(base_dataset) + ss_scale = load_ssa_benefit_projections(target_year) / max( + base_aggregates["weighted_ss_total"], + 1.0, + ) + earnings_scale = load_taxable_payroll_projections(target_year) / max( + base_aggregates["weighted_payroll_total"], + 1.0, + ) + pools = build_quantile_pools( + actual_summary, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + ) + candidates = generate_synthetic_candidates(pools) + exact_weights, solve_info = solve_synthetic_support(candidates, year=target_year) + exact_df = summarize_exact_candidates(candidates, exact_weights) + target_df = exact_df[exact_df["synthetic_weight"] > 0].head(top_n_targets).copy() + scaled_actual = build_scaled_actual_summary( + actual_summary, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + ) + + tax_unit_id_col = _period_column("person_tax_unit_id", base_year) + id_counters = { + entity_name: _next_entity_id(input_df[_period_column(columns[0], base_year)]) + for entity_name, columns in ENTITY_ID_COLUMNS.items() + } + id_counters["person"] = _next_entity_id(input_df[_period_column(PERSON_ID_COLUMN, base_year)]) + + clone_frames = [] + target_reports = [] + skipped_targets = [] + + for _, target_row in target_df.iterrows(): + target_candidate = candidates[int(target_row["candidate_idx"])] + donor_matches = match_real_donors_for_target( + target_row, + scaled_actual, + donors_per_target=donors_per_target, + ) + usable = donor_matches[donor_matches["distance"] <= max_distance_for_clone].copy() + if usable.empty: + skipped_targets.append( + { + "candidate_idx": int(target_row["candidate_idx"]), + "archetype": target_candidate.archetype, + "weight_share_pct": float(target_row["weight_share_pct"]), + "best_distance": float(donor_matches["distance"].min()), + } + ) + continue + + successful_clone_count = 0 + for _, donor_row in usable.iterrows(): + donor_rows = input_df[input_df[tax_unit_id_col] == int(donor_row["tax_unit_id"])].copy() + clone_df, id_counters = _clone_tax_unit_rows_to_target( + donor_rows, + base_year=base_year, + target_candidate=target_candidate, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + id_counters=id_counters, + clone_weight_scale=clone_weight_scale, + clone_weight_divisor=len(usable), + ) + if clone_df is None: + continue + clone_frames.append(clone_df) + successful_clone_count += 1 + target_reports.append( + { + "candidate_idx": int(target_row["candidate_idx"]), + "archetype": target_candidate.archetype, + "weight_share_pct": float(target_row["weight_share_pct"]), + "requested_donor_count": int(len(usable)), + "successful_clone_count": int(successful_clone_count), + } + ) + + augmented_df = ( + pd.concat([input_df, *clone_frames], ignore_index=True) + if clone_frames + else input_df.copy() + ) + report = { + "base_dataset": base_dataset, + "base_year": int(base_year), + "target_year": int(target_year), + "target_source": get_long_term_target_source(), + "solve_info": solve_info, + "top_n_targets": int(top_n_targets), + "donors_per_target": int(donors_per_target), + "max_distance_for_clone": float(max_distance_for_clone), + "clone_weight_scale": float(clone_weight_scale), + "base_household_count": int(input_df[_period_column("household_id", base_year)].nunique()), + "augmented_household_count": int(augmented_df[_period_column("household_id", base_year)].nunique()), + "base_person_count": int(len(input_df)), + "augmented_person_count": int(len(augmented_df)), + "target_reports": target_reports, + "skipped_targets": skipped_targets, + } + return augmented_df, report + + +def build_donor_backed_augmented_dataset( + *, + base_dataset: str, + base_year: int, + target_year: int, + top_n_targets: int = 20, + donors_per_target: int = 5, + max_distance_for_clone: float = 3.0, + clone_weight_scale: float = 0.1, +) -> tuple[Dataset, dict[str, object]]: + augmented_df, report = build_donor_backed_augmented_input_dataframe( + base_dataset=base_dataset, + base_year=base_year, + target_year=target_year, + top_n_targets=top_n_targets, + donors_per_target=donors_per_target, + max_distance_for_clone=max_distance_for_clone, + clone_weight_scale=clone_weight_scale, + ) + return Dataset.from_dataframe(augmented_df, base_year), report + + def _safe_split(numerator: float, denominator: float, fallback: float) -> float: if denominator <= 0: return fallback diff --git a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py index 4b7cf8b6d..f92b1253d 100644 --- a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py +++ b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py @@ -4,6 +4,7 @@ Usage: python run_household_projection.py [START_YEAR] [END_YEAR] [--profile PROFILE] [--target-source SOURCE] [--output-dir DIR] [--save-h5] [--allow-validation-failures] + python run_household_projection.py [START_YEAR] [END_YEAR] [--profile PROFILE] [--target-source SOURCE] [--support-augmentation-profile donor-backed-synthetic-v1] [--support-augmentation-target-year YEAR] python run_household_projection.py [START_YEAR] [END_YEAR] [--greg] [--use-ss] [--use-payroll] [--use-h6-reform] [--use-tob] [--save-h5] START_YEAR: Optional starting year (default: 2025) @@ -12,6 +13,8 @@ --target-source: Named long-term target source package --output-dir: Output directory for generated H5 files and metadata --allow-validation-failures: Record validation issues in metadata and continue instead of aborting the run + --support-augmentation-profile: Experimental late-year support expansion profile + --support-augmentation-target-year: Year whose extreme support is used to build the supplement --greg: Use GREG calibration instead of IPF (optional) --use-ss: Include Social Security benefit totals as calibration target (requires --greg) --use-payroll: Include taxable payroll totals as calibration target (requires --greg) @@ -22,6 +25,7 @@ Examples: python run_household_projection.py 2045 2045 --profile ss --target-source trustees_2025_current_law --save-h5 python run_household_projection.py 2025 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --save-h5 + python run_household_projection.py 2075 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --support-augmentation-profile donor-backed-synthetic-v1 --support-augmentation-target-year 2100 --allow-validation-failures """ import sys @@ -57,6 +61,9 @@ build_household_age_matrix, create_household_year_h5, ) +from prototype_synthetic_2100_support import ( + build_donor_backed_augmented_dataset, +) def create_h6_reform(): @@ -245,6 +252,8 @@ def create_h6_reform(): BASE_DATASET_PATH = DATASET_OPTIONS[SELECTED_DATASET]["path"] BASE_YEAR = DATASET_OPTIONS[SELECTED_DATASET]["base_year"] +SUPPORTED_AUGMENTATION_PROFILES = {"donor-backed-synthetic-v1"} + PROFILE_NAME = None if "--profile" in sys.argv: @@ -270,6 +279,80 @@ def create_h6_reform(): OUTPUT_DIR = sys.argv[output_dir_index + 1] del sys.argv[output_dir_index : output_dir_index + 2] +SUPPORT_AUGMENTATION_PROFILE = None +if "--support-augmentation-profile" in sys.argv: + augmentation_index = sys.argv.index("--support-augmentation-profile") + if augmentation_index + 1 >= len(sys.argv): + raise ValueError( + "--support-augmentation-profile requires a profile name" + ) + SUPPORT_AUGMENTATION_PROFILE = sys.argv[augmentation_index + 1] + del sys.argv[augmentation_index : augmentation_index + 2] + +SUPPORT_AUGMENTATION_TARGET_YEAR = None +if "--support-augmentation-target-year" in sys.argv: + target_year_index = sys.argv.index("--support-augmentation-target-year") + if target_year_index + 1 >= len(sys.argv): + raise ValueError( + "--support-augmentation-target-year requires a year" + ) + SUPPORT_AUGMENTATION_TARGET_YEAR = int(sys.argv[target_year_index + 1]) + del sys.argv[target_year_index : target_year_index + 2] + +SUPPORT_AUGMENTATION_START_YEAR = 2075 +if "--support-augmentation-start-year" in sys.argv: + start_year_index = sys.argv.index("--support-augmentation-start-year") + if start_year_index + 1 >= len(sys.argv): + raise ValueError( + "--support-augmentation-start-year requires a year" + ) + SUPPORT_AUGMENTATION_START_YEAR = int(sys.argv[start_year_index + 1]) + del sys.argv[start_year_index : start_year_index + 2] + +SUPPORT_AUGMENTATION_TOP_N_TARGETS = 20 +if "--support-augmentation-top-n-targets" in sys.argv: + top_n_index = sys.argv.index("--support-augmentation-top-n-targets") + if top_n_index + 1 >= len(sys.argv): + raise ValueError( + "--support-augmentation-top-n-targets requires an integer" + ) + SUPPORT_AUGMENTATION_TOP_N_TARGETS = int(sys.argv[top_n_index + 1]) + del sys.argv[top_n_index : top_n_index + 2] + +SUPPORT_AUGMENTATION_DONORS_PER_TARGET = 5 +if "--support-augmentation-donors-per-target" in sys.argv: + donor_index = sys.argv.index("--support-augmentation-donors-per-target") + if donor_index + 1 >= len(sys.argv): + raise ValueError( + "--support-augmentation-donors-per-target requires an integer" + ) + SUPPORT_AUGMENTATION_DONORS_PER_TARGET = int(sys.argv[donor_index + 1]) + del sys.argv[donor_index : donor_index + 2] + +SUPPORT_AUGMENTATION_MAX_DISTANCE = 3.0 +if "--support-augmentation-max-distance" in sys.argv: + distance_index = sys.argv.index("--support-augmentation-max-distance") + if distance_index + 1 >= len(sys.argv): + raise ValueError( + "--support-augmentation-max-distance requires a float" + ) + SUPPORT_AUGMENTATION_MAX_DISTANCE = float(sys.argv[distance_index + 1]) + del sys.argv[distance_index : distance_index + 2] + +SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE = 0.1 +if "--support-augmentation-clone-weight-scale" in sys.argv: + weight_scale_index = sys.argv.index( + "--support-augmentation-clone-weight-scale" + ) + if weight_scale_index + 1 >= len(sys.argv): + raise ValueError( + "--support-augmentation-clone-weight-scale requires a float" + ) + SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE = float( + sys.argv[weight_scale_index + 1] + ) + del sys.argv[weight_scale_index : weight_scale_index + 2] + ALLOW_VALIDATION_FAILURES = "--allow-validation-failures" in sys.argv if ALLOW_VALIDATION_FAILURES: sys.argv.remove("--allow-validation-failures") @@ -317,6 +400,22 @@ def create_h6_reform(): START_YEAR = int(sys.argv[1]) if len(sys.argv) > 1 else 2025 END_YEAR = int(sys.argv[2]) if len(sys.argv) > 2 else 2035 +if SUPPORT_AUGMENTATION_TARGET_YEAR is None: + SUPPORT_AUGMENTATION_TARGET_YEAR = END_YEAR + +if SUPPORT_AUGMENTATION_PROFILE is not None: + if SUPPORT_AUGMENTATION_PROFILE not in SUPPORTED_AUGMENTATION_PROFILES: + raise ValueError( + "Unsupported support augmentation profile: " + f"{SUPPORT_AUGMENTATION_PROFILE}" + ) + if START_YEAR < SUPPORT_AUGMENTATION_START_YEAR: + raise ValueError( + "Support augmentation is only supported for late-year runs. " + f"Received START_YEAR={START_YEAR}, requires >= " + f"{SUPPORT_AUGMENTATION_START_YEAR}." + ) + legacy_flags_used = any([USE_GREG, USE_SS, USE_PAYROLL, USE_H6_REFORM, USE_TOB]) if PROFILE_NAME and legacy_flags_used: raise ValueError("Use either --profile or legacy calibration flags, not both.") @@ -337,6 +436,9 @@ def create_h6_reform(): TARGET_SOURCE = get_long_term_target_source() TARGET_SOURCE_METADATA = describe_long_term_target_source(TARGET_SOURCE) +BASE_DATASET = BASE_DATASET_PATH +SUPPORT_AUGMENTATION_METADATA = None + CALIBRATION_METHOD = PROFILE.calibration_method USE_GREG = CALIBRATION_METHOD == "greg" USE_SS = PROFILE.use_ss @@ -374,6 +476,12 @@ def create_h6_reform(): print(f" Profile description: {PROFILE.description}") print(f" Target source: {TARGET_SOURCE}") print(f" Calibration method: {CALIBRATION_METHOD.upper()}") +if SUPPORT_AUGMENTATION_PROFILE: + print(f" Support augmentation: {SUPPORT_AUGMENTATION_PROFILE}") + print( + " Support augmentation target year: " + f"{SUPPORT_AUGMENTATION_TARGET_YEAR}" + ) if USE_SS: print(f" Including Social Security benefits constraint: Yes") if USE_PAYROLL: @@ -421,6 +529,64 @@ def create_h6_reform(): pop = target_matrix[:, idx].sum() print(f" {y}: {pop / 1e6:6.1f}M") +if SUPPORT_AUGMENTATION_PROFILE == "donor-backed-synthetic-v1": + print("\n" + "=" * 70) + print("STEP 1B: BUILD DONOR-BACKED LATE-YEAR SUPPORT") + print("=" * 70) + BASE_DATASET, augmentation_report = build_donor_backed_augmented_dataset( + base_dataset=BASE_DATASET_PATH, + base_year=BASE_YEAR, + target_year=SUPPORT_AUGMENTATION_TARGET_YEAR, + top_n_targets=SUPPORT_AUGMENTATION_TOP_N_TARGETS, + donors_per_target=SUPPORT_AUGMENTATION_DONORS_PER_TARGET, + max_distance_for_clone=SUPPORT_AUGMENTATION_MAX_DISTANCE, + clone_weight_scale=SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE, + ) + SUPPORT_AUGMENTATION_METADATA = { + "name": SUPPORT_AUGMENTATION_PROFILE, + "activation_start_year": SUPPORT_AUGMENTATION_START_YEAR, + "target_year": SUPPORT_AUGMENTATION_TARGET_YEAR, + "top_n_targets": SUPPORT_AUGMENTATION_TOP_N_TARGETS, + "donors_per_target": SUPPORT_AUGMENTATION_DONORS_PER_TARGET, + "max_distance_for_clone": SUPPORT_AUGMENTATION_MAX_DISTANCE, + "clone_weight_scale": SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE, + "report_summary": { + "base_household_count": augmentation_report["base_household_count"], + "augmented_household_count": augmentation_report[ + "augmented_household_count" + ], + "base_person_count": augmentation_report["base_person_count"], + "augmented_person_count": augmentation_report[ + "augmented_person_count" + ], + "successful_target_count": sum( + report["successful_clone_count"] > 0 + for report in augmentation_report["target_reports"] + ), + "skipped_target_count": len( + augmentation_report["skipped_targets"] + ), + }, + } + print( + " Base households -> augmented households: " + f"{augmentation_report['base_household_count']:,} -> " + f"{augmentation_report['augmented_household_count']:,}" + ) + print( + " Base people -> augmented people: " + f"{augmentation_report['base_person_count']:,} -> " + f"{augmentation_report['augmented_person_count']:,}" + ) + print( + " Successful target clones: " + f"{sum(report['successful_clone_count'] > 0 for report in augmentation_report['target_reports'])}" + ) + print( + " Skipped synthetic targets: " + f"{len(augmentation_report['skipped_targets'])}" + ) + # ========================================================================= # STEP 2: BUILD HOUSEHOLD AGE MATRIX # ========================================================================= @@ -428,7 +594,7 @@ def create_h6_reform(): print("STEP 2: BUILDING HOUSEHOLD AGE COMPOSITION") print("=" * 70) -sim = Microsimulation(dataset=BASE_DATASET_PATH) +sim = Microsimulation(dataset=BASE_DATASET) X, household_ids_unique, hh_id_to_idx = build_household_age_matrix(sim, n_ages) n_households = len(household_ids_unique) aggregated_age_cache: dict[int, tuple[np.ndarray, np.ndarray]] = {} @@ -467,7 +633,7 @@ def create_h6_reform(): for year_idx in range(n_years): year = START_YEAR + year_idx - sim = Microsimulation(dataset=BASE_DATASET_PATH) + sim = Microsimulation(dataset=BASE_DATASET) income_tax_hh = sim.calculate("income_tax", period=year, map_to="household") income_tax_baseline_total = income_tax_hh.sum() @@ -527,7 +693,7 @@ def create_h6_reform(): else: # Create and apply H6 reform h6_reform = create_h6_reform() - reform_sim = Microsimulation(dataset=BASE_DATASET_PATH, reform=h6_reform) + reform_sim = Microsimulation(dataset=BASE_DATASET, reform=h6_reform) # Calculate reform income tax income_tax_reform_hh = reform_sim.calculate( @@ -762,7 +928,7 @@ def create_h6_reform(): total_population[year_idx] = np.sum(y_target) if SAVE_H5: - h5_path = create_household_year_h5(year, w_new, BASE_DATASET_PATH, OUTPUT_DIR) + h5_path = create_household_year_h5(year, w_new, BASE_DATASET, OUTPUT_DIR) metadata_path = write_year_metadata( h5_path, year=year, @@ -770,6 +936,7 @@ def create_h6_reform(): profile=PROFILE.to_dict(), calibration_audit=calibration_audit, target_source=TARGET_SOURCE_METADATA, + support_augmentation=SUPPORT_AUGMENTATION_METADATA, ) update_dataset_manifest( OUTPUT_DIR, @@ -780,6 +947,7 @@ def create_h6_reform(): profile=PROFILE.to_dict(), calibration_audit=calibration_audit, target_source=TARGET_SOURCE_METADATA, + support_augmentation=SUPPORT_AUGMENTATION_METADATA, ) if year in display_years: print(f" Saved {year}.h5 and metadata") diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index 23a9c4671..ea09c6b84 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -948,3 +948,61 @@ def test_manifest_contains_invalid_artifacts_flag(tmp_path): assert manifest["contains_invalid_artifacts"] is True assert manifest["datasets"]["2030"]["validation_passed"] is True assert manifest["datasets"]["2091"]["validation_passed"] is False + + +def test_manifest_persists_support_augmentation_metadata(tmp_path): + profile = get_profile("ss-payroll-tob") + audit = { + "method_used": "entropy", + "fell_back_to_ipf": False, + "age_max_pct_error": 0.0, + "negative_weight_pct": 0.0, + "positive_weight_count": 70000, + "effective_sample_size": 5000.0, + "top_10_weight_share_pct": 1.5, + "top_100_weight_share_pct": 10.0, + "max_constraint_pct_error": 0.0, + "constraints": {}, + "validation_passed": True, + "validation_issues": [], + } + support_augmentation = { + "name": "donor-backed-synthetic-v1", + "activation_start_year": 2075, + "target_year": 2100, + "report_summary": { + "base_household_count": 41314, + "augmented_household_count": 41326, + }, + } + + year_2100 = tmp_path / "2100.h5" + year_2100.write_text("", encoding="utf-8") + metadata_path = write_year_metadata( + year_2100, + year=2100, + base_dataset_path="test.h5", + profile=profile.to_dict(), + calibration_audit=audit, + support_augmentation=support_augmentation, + ) + manifest_path = update_dataset_manifest( + tmp_path, + year=2100, + h5_path=year_2100, + metadata_path=metadata_path, + base_dataset_path="test.h5", + profile=profile.to_dict(), + calibration_audit=audit, + support_augmentation=support_augmentation, + ) + + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + assert metadata["support_augmentation"]["name"] == "donor-backed-synthetic-v1" + assert ( + manifest["support_augmentation"]["report_summary"][ + "augmented_household_count" + ] + == 41326 + ) From d93eed873c8e1be1be7ba605cfa620fef317bdd0 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 1 Apr 2026 10:48:29 -0400 Subject: [PATCH 23/31] Add structural donor composite late-tail prototype --- .../datasets/cps/long_term/README.md | 12 +- .../prototype_synthetic_2100_support.py | 769 +++++++++++++++++- .../cps/long_term/run_household_projection.py | 41 +- .../test_long_term_calibration_contract.py | 83 ++ 4 files changed, 877 insertions(+), 28 deletions(-) diff --git a/policyengine_us_data/datasets/cps/long_term/README.md b/policyengine_us_data/datasets/cps/long_term/README.md index dc361120e..8f1360d9e 100644 --- a/policyengine_us_data/datasets/cps/long_term/README.md +++ b/policyengine_us_data/datasets/cps/long_term/README.md @@ -12,6 +12,9 @@ python run_household_projection.py 2100 --profile ss-payroll-tob --target-source # Experimental: donor-backed late-year support augmentation for tail-year runs python run_household_projection.py 2075 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --support-augmentation-profile donor-backed-synthetic-v1 --support-augmentation-target-year 2100 --allow-validation-failures +# Experimental: role-based donor composites assembled into late-year support +python run_household_projection.py 2075 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --support-augmentation-profile donor-backed-composite-v1 --support-augmentation-target-year 2100 --allow-validation-failures + # IPF with only age distribution constraints (faster, less accurate) python run_household_projection.py 2050 --profile age-only @@ -24,7 +27,7 @@ python run_household_projection.py 2100 --profile ss - `--profile`: Named calibration contract. Recommended over legacy flags. - `--target-source`: Named long-term target source package. - `--output-dir`: Output directory for generated H5 files and metadata sidecars. -- `--support-augmentation-profile`: Experimental late-year support expansion mode. Currently supports `donor-backed-synthetic-v1`. +- `--support-augmentation-profile`: Experimental late-year support expansion mode. Currently supports `donor-backed-synthetic-v1` and `donor-backed-composite-v1`. - `--support-augmentation-target-year`: Extreme year used to build the donor-backed supplement (defaults to `END_YEAR`). - `--support-augmentation-start-year`: Earliest run year allowed for augmentation (defaults to `2075`). - `--support-augmentation-top-n-targets`: Number of dominant synthetic target types to map back to real donors (default `20`). @@ -79,6 +82,13 @@ python run_household_projection.py 2100 --profile ss - Intended to test whether donor-backed synthetic support improves late-year microsim feasibility without resorting to fully free synthetic records - Current status: integrated into the runner and fully auditable in metadata, but still diagnostic. The first `2100` end-to-end run did not materially improve the late-tail calibration frontier or support-concentration metrics. +**Role-Based Donor Composites** +- Experimental structural extension of the donor-backed approach +- Recombines older-beneficiary donors, payroll-rich worker donors, and dependent structure into synthetic household candidates before assembling actual augmented rows +- In the synthetic support lab, this materially improves the `2100` exact-fit basis: exact entropy fit with `360` positive candidates, ESS about `95.8`, and top-10 weight share about `25.6%` +- The actual-row augmented dataset builder is now available in the runner as `donor-backed-composite-v1` +- Current runner status: the first full `2100` end-to-end run with the default composite supplement added `270` structural clones and modestly improved support concentration (`ESS 11.4 -> 13.3`, top-10 share `84.9% -> 79.4%`, lower TOB overshoot), but it did not yet move the main `SS + payroll` frontier off the current `-33.8% / -35.0%` late-tail bound + **GREG (Generalized Regression Estimator)** - Legacy linear calibration path retained for explicit flag-based runs - Can hit constraints exactly, but may produce negative weights in far-horizon years diff --git a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py index 12bcec437..87fa6b8af 100644 --- a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py +++ b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py @@ -10,20 +10,36 @@ from policyengine_core.data.dataset import Dataset from policyengine_us import Microsimulation -from calibration import ( - assess_nonnegative_feasibility, - calibrate_entropy, - calibrate_entropy_bounded, - densify_lp_solution, -) -from projection_utils import aggregate_age_targets, build_age_bins -from ssa_data import ( - get_long_term_target_source, - load_ssa_age_projections, - load_ssa_benefit_projections, - load_taxable_payroll_projections, - set_long_term_target_source, -) +try: + from .calibration import ( + assess_nonnegative_feasibility, + calibrate_entropy, + calibrate_entropy_bounded, + densify_lp_solution, + ) + from .projection_utils import aggregate_age_targets, build_age_bins + from .ssa_data import ( + get_long_term_target_source, + load_ssa_age_projections, + load_ssa_benefit_projections, + load_taxable_payroll_projections, + set_long_term_target_source, + ) +except ImportError: # pragma: no cover - script execution fallback + from calibration import ( + assess_nonnegative_feasibility, + calibrate_entropy, + calibrate_entropy_bounded, + densify_lp_solution, + ) + from projection_utils import aggregate_age_targets, build_age_bins + from ssa_data import ( + get_long_term_target_source, + load_ssa_age_projections, + load_ssa_benefit_projections, + load_taxable_payroll_projections, + set_long_term_target_source, + ) DEFAULT_DATASET = "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5" @@ -713,6 +729,7 @@ def build_synthetic_constraint_problem( candidates: list[SyntheticCandidate], *, year: int, + baseline_weights: np.ndarray | None = None, ) -> dict[str, object]: age_targets = load_ssa_age_projections(start_year=year, end_year=year) age_bins = build_age_bins(n_ages=age_targets.shape[0], bucket_size=5) @@ -723,6 +740,12 @@ def build_synthetic_constraint_problem( [candidate.payroll_total for candidate in candidates], dtype=float, ) + if baseline_weights is None: + baseline_weights = np.ones(len(candidates), dtype=float) + else: + baseline_weights = np.asarray(baseline_weights, dtype=float) + if len(baseline_weights) != len(candidates): + raise ValueError("baseline_weights must align with candidates") return { "age_bins": age_bins, "aggregated_age_targets": aggregated_age_targets, @@ -731,7 +754,7 @@ def build_synthetic_constraint_problem( "ss_target": float(load_ssa_benefit_projections(year)), "payroll_values": payroll_values, "payroll_target": float(load_taxable_payroll_projections(year)), - "baseline_weights": np.ones(len(candidates), dtype=float), + "baseline_weights": baseline_weights, } @@ -770,6 +793,316 @@ def build_scaled_actual_summary( return scaled +def _target_head_payroll_share(candidate: SyntheticCandidate) -> float: + return _safe_split( + candidate.head_wages, + candidate.payroll_total, + 1.0 if candidate.spouse_age is None else 0.5, + ) + + +def _target_head_ss_share(candidate: SyntheticCandidate) -> float: + return _safe_split( + candidate.head_ss, + candidate.ss_total, + 1.0 if candidate.spouse_age is None else 0.5, + ) + + +def _target_worker_age(candidate: SyntheticCandidate) -> float: + if candidate.spouse_age is not None and candidate.spouse_wages > candidate.head_wages: + return float(candidate.spouse_age) + return float(candidate.head_age) + + +def match_older_role_donors( + target_candidate: SyntheticCandidate, + scaled_actual_summary: pd.DataFrame, + *, + donors_per_target: int, +) -> pd.DataFrame: + subset = scaled_actual_summary[ + (scaled_actual_summary["ss_total"] > 0) + & (scaled_actual_summary["head_age"] >= 55) + ].copy() + if subset.empty: + return subset + target_spouse_age = ( + -1 + if target_candidate.spouse_age is None or target_candidate.spouse_age < 65 + else target_candidate.spouse_age + ) + target_head_ss_share = _target_head_ss_share(target_candidate) + donor_head_ss_share = _safe_series_split( + subset["scaled_head_ss"], + subset["scaled_ss_total"], + target_head_ss_share, + ) + subset["distance"] = ( + (subset["head_age"] - target_candidate.head_age).abs() / 5.0 + + (subset["spouse_age_filled"] - target_spouse_age).abs() / 7.5 + + np.abs( + np.log1p(subset["scaled_ss_total"]) - np.log1p(target_candidate.ss_total) + ) + + 0.5 + * np.abs(donor_head_ss_share - target_head_ss_share) + + 0.15 + * subset["archetype"] + .isin( + { + "older_beneficiary_single", + "older_beneficiary_couple", + "older_worker_single", + "older_worker_couple", + "mixed_retiree_worker_couple", + "older_plus_prime_worker_family", + } + ) + .rsub(1) + .astype(float) + ) + return subset.nsmallest(donors_per_target, "distance").copy() + + +def match_worker_role_donors( + target_candidate: SyntheticCandidate, + scaled_actual_summary: pd.DataFrame, + *, + donors_per_target: int, +) -> pd.DataFrame: + subset = scaled_actual_summary[ + scaled_actual_summary["payroll_total"] > 0 + ].copy() + target_dependent_count = len(target_candidate.dependent_ages) + target_spouse_present = target_candidate.spouse_age is not None + if target_dependent_count > 0: + family_subset = subset[subset["dependent_count"] == target_dependent_count] + if not family_subset.empty: + subset = family_subset.copy() + if target_spouse_present: + spouse_subset = subset[subset["spouse_present"]] + if not spouse_subset.empty: + subset = spouse_subset.copy() + target_worker_age = _target_worker_age(target_candidate) + target_head_payroll_share = _target_head_payroll_share(target_candidate) + donor_head_payroll_share = _safe_series_split( + subset["scaled_head_payroll"], + subset["scaled_payroll_total"], + target_head_payroll_share, + ) + subset["distance"] = ( + (subset["head_age"] - target_worker_age).abs() / 5.0 + + (subset["dependent_count"] - target_dependent_count).abs() * 0.75 + + np.abs( + np.log1p(subset["scaled_payroll_total"]) + - np.log1p(target_candidate.payroll_total) + ) + + 0.5 + * np.abs(donor_head_payroll_share - target_head_payroll_share) + + 0.25 * subset["spouse_present"].ne(target_spouse_present).astype(float) + ) + return subset.nsmallest(donors_per_target, "distance").copy() + + +def build_role_donor_composite_candidate( + target_candidate: SyntheticCandidate, + *, + older_donor_row: pd.Series | None, + worker_donor_row: pd.Series | None, + earnings_scale: float, +) -> SyntheticCandidate: + target_head_payroll_share = _target_head_payroll_share(target_candidate) + target_head_ss_share = _target_head_ss_share(target_candidate) + if ( + worker_donor_row is not None + and target_candidate.head_wages > 0 + and target_candidate.spouse_wages > 0 + ): + head_payroll_share = _safe_split( + float(worker_donor_row["scaled_head_payroll"]), + float(worker_donor_row["scaled_payroll_total"]), + target_head_payroll_share, + ) + else: + head_payroll_share = target_head_payroll_share + if ( + older_donor_row is not None + and target_candidate.head_ss > 0 + and target_candidate.spouse_ss > 0 + ): + head_ss_share = _safe_split( + float(older_donor_row["scaled_head_ss"]), + float(older_donor_row["scaled_ss_total"]), + target_head_ss_share, + ) + else: + head_ss_share = target_head_ss_share + + payroll_total = target_candidate.payroll_total + ss_total = target_candidate.ss_total + pension_income = 0.0 + dividend_income = 0.0 + if older_donor_row is not None: + pension_income += float(older_donor_row["pension_income"]) * earnings_scale + dividend_income += float(older_donor_row["dividend_income"]) * earnings_scale + if worker_donor_row is not None and target_candidate.ss_total <= 0: + pension_income += float(worker_donor_row["pension_income"]) * earnings_scale + dividend_income += float(worker_donor_row["dividend_income"]) * earnings_scale + + return SyntheticCandidate( + archetype=f"{target_candidate.archetype}_role_donor", + head_age=target_candidate.head_age, + spouse_age=target_candidate.spouse_age, + dependent_ages=target_candidate.dependent_ages, + head_wages=payroll_total * head_payroll_share, + spouse_wages=payroll_total * (1.0 - head_payroll_share), + head_ss=ss_total * head_ss_share, + spouse_ss=ss_total * (1.0 - head_ss_share), + pension_income=pension_income, + dividend_income=dividend_income, + ) + + +def build_role_donor_composites( + candidates: list[SyntheticCandidate], + weights: np.ndarray, + actual_summary: pd.DataFrame, + *, + ss_scale: float, + earnings_scale: float, + top_n_targets: int, + older_donors_per_target: int, + worker_donors_per_target: int, + max_older_distance: float = 3.0, + max_worker_distance: float = 3.0, +) -> tuple[list[SyntheticCandidate], np.ndarray, dict[str, object]]: + exact_df = summarize_exact_candidates(candidates, weights) + target_df = exact_df[exact_df["synthetic_weight"] > 0].head(top_n_targets).copy() + scaled_actual = build_scaled_actual_summary( + actual_summary, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + ) + + composite_candidates: list[SyntheticCandidate] = [] + composite_weights: list[float] = [] + composite_records: list[dict[str, object]] = [] + skipped_targets: list[dict[str, object]] = [] + + total_weight = max(float(weights.sum()), 1.0) + + for _, target_row in target_df.iterrows(): + target_candidate = candidates[int(target_row["candidate_idx"])] + older_donors = [None] + if target_candidate.ss_total > 0: + matched = match_older_role_donors( + target_candidate, + scaled_actual, + donors_per_target=older_donors_per_target, + ) + usable = matched[matched["distance"] <= max_older_distance].copy() + if usable.empty: + skipped_targets.append( + { + "candidate_idx": int(target_row["candidate_idx"]), + "archetype": target_candidate.archetype, + "reason": "no_older_donor", + "best_distance": float(matched["distance"].min()) + if not matched.empty + else None, + } + ) + continue + older_donors = [row for _, row in usable.iterrows()] + + worker_donors = [None] + if target_candidate.payroll_total > 0: + matched = match_worker_role_donors( + target_candidate, + scaled_actual, + donors_per_target=worker_donors_per_target, + ) + usable = matched[matched["distance"] <= max_worker_distance].copy() + if usable.empty: + skipped_targets.append( + { + "candidate_idx": int(target_row["candidate_idx"]), + "archetype": target_candidate.archetype, + "reason": "no_worker_donor", + "best_distance": float(matched["distance"].min()) + if not matched.empty + else None, + } + ) + continue + worker_donors = [row for _, row in usable.iterrows()] + + target_weight = float(target_row["synthetic_weight"]) + combination_count = max(len(older_donors) * len(worker_donors), 1) + per_candidate_weight = target_weight / combination_count + for older_donor in older_donors: + for worker_donor in worker_donors: + composite_candidates.append( + build_role_donor_composite_candidate( + target_candidate, + older_donor_row=older_donor, + worker_donor_row=worker_donor, + earnings_scale=earnings_scale, + ) + ) + composite_weights.append(per_candidate_weight) + composite_records.append( + { + "candidate_idx": int(target_row["candidate_idx"]), + "archetype": target_candidate.archetype, + "older_tax_unit_id": ( + None + if older_donor is None + else int(older_donor["tax_unit_id"]) + ), + "worker_tax_unit_id": ( + None + if worker_donor is None + else int(worker_donor["tax_unit_id"]) + ), + "older_distance": ( + None + if older_donor is None + else float(older_donor["distance"]) + ), + "worker_distance": ( + None + if worker_donor is None + else float(worker_donor["distance"]) + ), + "assigned_weight_share_pct": float( + per_candidate_weight / total_weight * 100 + ), + } + ) + + prior_weights = np.asarray(composite_weights, dtype=float) + probe_summary = summarize_solution( + composite_candidates, + prior_weights, + actual_summary, + ) + return ( + composite_candidates, + prior_weights, + { + "top_n_targets": int(top_n_targets), + "older_donors_per_target": int(older_donors_per_target), + "worker_donors_per_target": int(worker_donors_per_target), + "max_older_distance": float(max_older_distance), + "max_worker_distance": float(max_worker_distance), + "skipped_targets": skipped_targets, + "composite_records": composite_records[:100], + "prior_summary": probe_summary, + }, + ) + + def summarize_exact_candidates( candidates: list[SyntheticCandidate], weights: np.ndarray, @@ -1097,6 +1430,172 @@ def _clone_tax_unit_rows_to_target( return cloned, id_counters +def _compose_role_donor_rows_to_target( + older_donor_rows: pd.DataFrame | None, + worker_donor_rows: pd.DataFrame | None, + *, + base_year: int, + target_candidate: SyntheticCandidate, + ss_scale: float, + earnings_scale: float, + id_counters: dict[str, int], + clone_weight_scale: float, + clone_weight_divisor: int, +) -> tuple[pd.DataFrame, dict[str, int]] | tuple[None, dict[str, int]]: + age_col = _period_column("age", base_year) + household_weight_col = _period_column("household_weight", base_year) + person_weight_col = _period_column("person_weight", base_year) + person_id_col = _period_column(PERSON_ID_COLUMN, base_year) + + def _adult_rows(df: pd.DataFrame | None) -> pd.DataFrame: + if df is None: + return pd.DataFrame(columns=[] if df is None else df.columns) + return df[df[age_col] >= 18].sort_values(age_col, ascending=False) + + def _dependent_rows(df: pd.DataFrame | None) -> pd.DataFrame: + if df is None: + return pd.DataFrame(columns=[] if df is None else df.columns) + return df[df[age_col] < 18].sort_values(age_col, ascending=False) + + older_adults = _adult_rows(older_donor_rows) + worker_adults = _adult_rows(worker_donor_rows) + worker_dependents = _dependent_rows(worker_donor_rows) + + selected_rows: list[pd.Series] = [] + head_target_older = target_candidate.head_ss > 0 or target_candidate.head_age >= 65 + head_source_rows = older_adults if head_target_older and not older_adults.empty else worker_adults + if head_source_rows.empty: + return None, id_counters + head_row = head_source_rows.iloc[0].copy() + selected_rows.append(head_row) + + spouse_row = None + if target_candidate.spouse_age is not None: + if target_candidate.spouse_age >= 65 and len(older_adults) >= 2: + spouse_row = older_adults.iloc[1].copy() + elif not worker_adults.empty: + worker_candidates = worker_adults.iloc[1:] if worker_adults.index[0] == head_row.name else worker_adults + if worker_candidates.empty: + worker_candidates = worker_adults + spouse_idx = (worker_candidates[age_col] - target_candidate.spouse_age).abs().idxmin() + spouse_row = worker_candidates.loc[spouse_idx].copy() + elif len(older_adults) >= 2: + spouse_row = older_adults.iloc[1].copy() + if spouse_row is None: + return None, id_counters + selected_rows.append(spouse_row) + + if len(target_candidate.dependent_ages) > 0: + if len(worker_dependents) < len(target_candidate.dependent_ages): + return None, id_counters + selected_rows.extend( + [ + row.copy() + for _, row in worker_dependents.iloc[ + : len(target_candidate.dependent_ages) + ].iterrows() + ] + ) + + cloned = pd.DataFrame(selected_rows).copy() + household_id = id_counters["household"] + id_counters["household"] += 1 + for entity_name, columns in ENTITY_ID_COLUMNS.items(): + entity_id = id_counters[entity_name] + id_counters[entity_name] += 1 + for raw_column in columns: + column = _period_column(raw_column, base_year) + if column in cloned.columns: + cloned[column] = entity_id if entity_name != "household" else household_id + cloned[_period_column("household_id", base_year)] = household_id + cloned[_period_column("person_household_id", base_year)] = household_id + + person_ids = range(id_counters["person"], id_counters["person"] + len(cloned)) + id_counters["person"] += len(cloned) + cloned[person_id_col] = _cast_mapped_ids( + cloned[person_id_col], + pd.Series(list(person_ids), index=cloned.index), + ) + + if household_weight_col in cloned.columns: + cloned[household_weight_col] = ( + cloned[household_weight_col].astype(float) + * clone_weight_scale + / max(clone_weight_divisor, 1) + ) + if person_weight_col in cloned.columns: + cloned[person_weight_col] = ( + cloned[person_weight_col].astype(float) + * clone_weight_scale + / max(clone_weight_divisor, 1) + ) + + head_idx = cloned.index[0] + spouse_idx = cloned.index[1] if target_candidate.spouse_age is not None else None + dependent_indices = ( + cloned.index[2 : 2 + len(target_candidate.dependent_ages)] + if target_candidate.spouse_age is not None + else cloned.index[1 : 1 + len(target_candidate.dependent_ages)] + ) + + cloned.loc[head_idx, age_col] = float(target_candidate.head_age) + if spouse_idx is not None: + cloned.loc[spouse_idx, age_col] = float(target_candidate.spouse_age) + for dep_idx, dep_age in zip(dependent_indices, target_candidate.dependent_ages): + cloned.loc[dep_idx, age_col] = float(dep_age) + + payroll_columns = tuple(_period_column(component, base_year) for component in PAYROLL_COMPONENTS) + ss_columns = tuple(_period_column(component, base_year) for component in SS_COMPONENTS) + qbi_col = _period_column("w2_wages_from_qualified_business", base_year) + + target_head_payroll = float(target_candidate.head_wages) / max(earnings_scale, 1e-12) + target_spouse_payroll = float(target_candidate.spouse_wages) / max(earnings_scale, 1e-12) + target_head_ss = float(target_candidate.head_ss) / max(ss_scale, 1e-12) + target_spouse_ss = float(target_candidate.spouse_ss) / max(ss_scale, 1e-12) + + cloned.loc[head_idx] = _scale_person_components( + cloned.loc[head_idx].copy(), + payroll_columns, + target_head_payroll, + ) + cloned.loc[head_idx] = _scale_person_components( + cloned.loc[head_idx].copy(), + ss_columns, + target_head_ss, + ) + if spouse_idx is not None: + cloned.loc[spouse_idx] = _scale_person_components( + cloned.loc[spouse_idx].copy(), + payroll_columns, + target_spouse_payroll, + ) + cloned.loc[spouse_idx] = _scale_person_components( + cloned.loc[spouse_idx].copy(), + ss_columns, + target_spouse_ss, + ) + for dep_idx in dependent_indices: + cloned.loc[dep_idx] = _scale_person_components( + cloned.loc[dep_idx].copy(), + payroll_columns, + 0.0, + ) + cloned.loc[dep_idx] = _scale_person_components( + cloned.loc[dep_idx].copy(), + ss_columns, + 0.0, + ) + if qbi_col in cloned.columns: + cloned.loc[dep_idx, qbi_col] = 0.0 + + if qbi_col in cloned.columns and head_idx in cloned.index: + cloned.loc[head_idx, qbi_col] = 0.0 + if spouse_idx is not None: + cloned.loc[spouse_idx, qbi_col] = 0.0 + + return cloned, id_counters + + def build_donor_backed_augmented_input_dataframe( *, base_dataset: str, @@ -1216,6 +1715,191 @@ def build_donor_backed_augmented_input_dataframe( return augmented_df, report +def build_role_composite_augmented_input_dataframe( + *, + base_dataset: str, + base_year: int, + target_year: int, + top_n_targets: int = 20, + donors_per_target: int = 5, + max_older_distance: float = 3.0, + max_worker_distance: float = 3.0, + clone_weight_scale: float = 0.1, +) -> tuple[pd.DataFrame, dict[str, object]]: + sim = Microsimulation(dataset=base_dataset) + input_df = sim.to_input_dataframe() + actual_summary = build_actual_tax_unit_summary(base_dataset) + base_aggregates = load_base_aggregates(base_dataset) + ss_scale = load_ssa_benefit_projections(target_year) / max( + base_aggregates["weighted_ss_total"], + 1.0, + ) + earnings_scale = load_taxable_payroll_projections(target_year) / max( + base_aggregates["weighted_payroll_total"], + 1.0, + ) + pools = build_quantile_pools( + actual_summary, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + ) + candidates = generate_synthetic_candidates(pools) + exact_weights, solve_info = solve_synthetic_support(candidates, year=target_year) + exact_df = summarize_exact_candidates(candidates, exact_weights) + target_df = exact_df[exact_df["synthetic_weight"] > 0].head(top_n_targets).copy() + scaled_actual = build_scaled_actual_summary( + actual_summary, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + ) + + tax_unit_id_col = _period_column("person_tax_unit_id", base_year) + id_counters = { + entity_name: _next_entity_id(input_df[_period_column(columns[0], base_year)]) + for entity_name, columns in ENTITY_ID_COLUMNS.items() + } + id_counters["person"] = _next_entity_id( + input_df[_period_column(PERSON_ID_COLUMN, base_year)] + ) + + clone_frames = [] + target_reports = [] + skipped_targets = [] + + for _, target_row in target_df.iterrows(): + target_candidate = candidates[int(target_row["candidate_idx"])] + older_matches = ( + match_older_role_donors( + target_candidate, + scaled_actual, + donors_per_target=donors_per_target, + ) + if target_candidate.ss_total > 0 + else pd.DataFrame([{}]) + ) + worker_matches = ( + match_worker_role_donors( + target_candidate, + scaled_actual, + donors_per_target=donors_per_target, + ) + if target_candidate.payroll_total > 0 + else pd.DataFrame([{}]) + ) + if target_candidate.ss_total > 0: + older_matches = older_matches[ + older_matches["distance"] <= max_older_distance + ].copy() + if target_candidate.payroll_total > 0: + worker_matches = worker_matches[ + worker_matches["distance"] <= max_worker_distance + ].copy() + if older_matches.empty or worker_matches.empty: + skipped_targets.append( + { + "candidate_idx": int(target_row["candidate_idx"]), + "archetype": target_candidate.archetype, + "weight_share_pct": float(target_row["weight_share_pct"]), + "older_match_count": int(len(older_matches)), + "worker_match_count": int(len(worker_matches)), + } + ) + continue + + successful_clone_count = 0 + combination_count = len(older_matches) * len(worker_matches) + for _, older_row in older_matches.iterrows(): + older_rows = ( + input_df[input_df[tax_unit_id_col] == int(older_row["tax_unit_id"])].copy() + if target_candidate.ss_total > 0 + else None + ) + for _, worker_row in worker_matches.iterrows(): + worker_rows = ( + input_df[input_df[tax_unit_id_col] == int(worker_row["tax_unit_id"])].copy() + if target_candidate.payroll_total > 0 + else None + ) + clone_df, id_counters = _compose_role_donor_rows_to_target( + older_rows, + worker_rows, + base_year=base_year, + target_candidate=target_candidate, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + id_counters=id_counters, + clone_weight_scale=clone_weight_scale, + clone_weight_divisor=combination_count, + ) + if clone_df is None: + continue + clone_frames.append(clone_df) + successful_clone_count += 1 + target_reports.append( + { + "candidate_idx": int(target_row["candidate_idx"]), + "archetype": target_candidate.archetype, + "weight_share_pct": float(target_row["weight_share_pct"]), + "older_match_count": int(len(older_matches)), + "worker_match_count": int(len(worker_matches)), + "successful_clone_count": int(successful_clone_count), + } + ) + + augmented_df = ( + pd.concat([input_df, *clone_frames], ignore_index=True) + if clone_frames + else input_df.copy() + ) + report = { + "base_dataset": base_dataset, + "base_year": int(base_year), + "target_year": int(target_year), + "target_source": get_long_term_target_source(), + "solve_info": solve_info, + "top_n_targets": int(top_n_targets), + "donors_per_target": int(donors_per_target), + "max_older_distance": float(max_older_distance), + "max_worker_distance": float(max_worker_distance), + "clone_weight_scale": float(clone_weight_scale), + "base_household_count": int( + input_df[_period_column("household_id", base_year)].nunique() + ), + "augmented_household_count": int( + augmented_df[_period_column("household_id", base_year)].nunique() + ), + "base_person_count": int(len(input_df)), + "augmented_person_count": int(len(augmented_df)), + "target_reports": target_reports, + "skipped_targets": skipped_targets, + } + return augmented_df, report + + +def build_role_composite_augmented_dataset( + *, + base_dataset: str, + base_year: int, + target_year: int, + top_n_targets: int = 20, + donors_per_target: int = 5, + max_older_distance: float = 3.0, + max_worker_distance: float = 3.0, + clone_weight_scale: float = 0.1, +) -> tuple[Dataset, dict[str, object]]: + augmented_df, report = build_role_composite_augmented_input_dataframe( + base_dataset=base_dataset, + base_year=base_year, + target_year=target_year, + top_n_targets=top_n_targets, + donors_per_target=donors_per_target, + max_older_distance=max_older_distance, + max_worker_distance=max_worker_distance, + clone_weight_scale=clone_weight_scale, + ) + return Dataset.from_dataframe(augmented_df, base_year), report + + def build_donor_backed_augmented_dataset( *, base_dataset: str, @@ -1244,6 +1928,19 @@ def _safe_split(numerator: float, denominator: float, fallback: float) -> float: return float(numerator / denominator) +def _safe_series_split( + numerator: pd.Series, + denominator: pd.Series, + fallback: float, +) -> pd.Series: + numerator = numerator.astype(float) + denominator = denominator.astype(float) + result = pd.Series(fallback, index=numerator.index, dtype=float) + positive = denominator > 0 + result.loc[positive] = numerator.loc[positive] / denominator.loc[positive] + return result + + def build_donor_backed_clones( candidates: list[SyntheticCandidate], weights: np.ndarray, @@ -1366,8 +2063,13 @@ def solve_synthetic_support( year: int, max_constraint_error_pct: float = 0.0, warm_weights: np.ndarray | None = None, + baseline_weights: np.ndarray | None = None, ) -> tuple[np.ndarray, dict[str, object]]: - problem = build_synthetic_constraint_problem(candidates, year=year) + problem = build_synthetic_constraint_problem( + candidates, + year=year, + baseline_weights=baseline_weights, + ) X = problem["X"] aggregated_age_targets = problem["aggregated_age_targets"] ss_values = problem["ss_values"] @@ -1737,6 +2439,40 @@ def main() -> int: top_n_targets=args.donor_probe_top_n, donors_per_target=args.donor_probe_k, ) + ( + role_composite_candidates, + role_composite_prior, + role_donor_composite_probe, + ) = build_role_donor_composites( + candidates, + weights, + actual_summary, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + top_n_targets=args.donor_probe_top_n, + older_donors_per_target=args.donor_probe_k, + worker_donors_per_target=args.donor_probe_k, + ) + role_donor_composite_result: dict[str, object] = { + "candidate_count": int(len(role_composite_candidates)), + "prior_summary": role_donor_composite_probe["prior_summary"], + "skipped_targets": role_donor_composite_probe["skipped_targets"], + "composite_records": role_donor_composite_probe["composite_records"], + } + if role_composite_candidates: + role_weights, role_solve_info = solve_synthetic_support( + role_composite_candidates, + year=args.year, + baseline_weights=role_composite_prior, + ) + role_donor_composite_result["solve_info"] = role_solve_info + role_donor_composite_result["solution_summary"] = summarize_solution( + role_composite_candidates, + role_weights, + actual_summary, + ) + else: + role_donor_composite_result["solve_info"] = {"status": "no_candidates"} epsilon_path_results = [] for epsilon in parse_epsilon_path(args.epsilon_path): epsilon_weights, epsilon_solve_info = solve_synthetic_support( @@ -1780,6 +2516,7 @@ def main() -> int: "synthetic_solution": solution_summary, "donor_probe": donor_probe, "donor_backed_clone_probe": donor_backed_clone_probe, + "role_donor_composite_probe": role_donor_composite_result, "epsilon_path": epsilon_path_results, } diff --git a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py index f92b1253d..a14614b8f 100644 --- a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py +++ b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py @@ -63,6 +63,7 @@ ) from prototype_synthetic_2100_support import ( build_donor_backed_augmented_dataset, + build_role_composite_augmented_dataset, ) @@ -252,7 +253,10 @@ def create_h6_reform(): BASE_DATASET_PATH = DATASET_OPTIONS[SELECTED_DATASET]["path"] BASE_YEAR = DATASET_OPTIONS[SELECTED_DATASET]["base_year"] -SUPPORTED_AUGMENTATION_PROFILES = {"donor-backed-synthetic-v1"} +SUPPORTED_AUGMENTATION_PROFILES = { + "donor-backed-synthetic-v1", + "donor-backed-composite-v1", +} PROFILE_NAME = None @@ -529,19 +533,34 @@ def create_h6_reform(): pop = target_matrix[:, idx].sum() print(f" {y}: {pop / 1e6:6.1f}M") -if SUPPORT_AUGMENTATION_PROFILE == "donor-backed-synthetic-v1": +if SUPPORT_AUGMENTATION_PROFILE in { + "donor-backed-synthetic-v1", + "donor-backed-composite-v1", +}: print("\n" + "=" * 70) print("STEP 1B: BUILD DONOR-BACKED LATE-YEAR SUPPORT") print("=" * 70) - BASE_DATASET, augmentation_report = build_donor_backed_augmented_dataset( - base_dataset=BASE_DATASET_PATH, - base_year=BASE_YEAR, - target_year=SUPPORT_AUGMENTATION_TARGET_YEAR, - top_n_targets=SUPPORT_AUGMENTATION_TOP_N_TARGETS, - donors_per_target=SUPPORT_AUGMENTATION_DONORS_PER_TARGET, - max_distance_for_clone=SUPPORT_AUGMENTATION_MAX_DISTANCE, - clone_weight_scale=SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE, - ) + if SUPPORT_AUGMENTATION_PROFILE == "donor-backed-synthetic-v1": + BASE_DATASET, augmentation_report = build_donor_backed_augmented_dataset( + base_dataset=BASE_DATASET_PATH, + base_year=BASE_YEAR, + target_year=SUPPORT_AUGMENTATION_TARGET_YEAR, + top_n_targets=SUPPORT_AUGMENTATION_TOP_N_TARGETS, + donors_per_target=SUPPORT_AUGMENTATION_DONORS_PER_TARGET, + max_distance_for_clone=SUPPORT_AUGMENTATION_MAX_DISTANCE, + clone_weight_scale=SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE, + ) + else: + BASE_DATASET, augmentation_report = build_role_composite_augmented_dataset( + base_dataset=BASE_DATASET_PATH, + base_year=BASE_YEAR, + target_year=SUPPORT_AUGMENTATION_TARGET_YEAR, + top_n_targets=SUPPORT_AUGMENTATION_TOP_N_TARGETS, + donors_per_target=SUPPORT_AUGMENTATION_DONORS_PER_TARGET, + max_older_distance=SUPPORT_AUGMENTATION_MAX_DISTANCE, + max_worker_distance=SUPPORT_AUGMENTATION_MAX_DISTANCE, + clone_weight_scale=SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE, + ) SUPPORT_AUGMENTATION_METADATA = { "name": SUPPORT_AUGMENTATION_PROFILE, "activation_start_year": SUPPORT_AUGMENTATION_START_YEAR, diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index ea09c6b84..50e278a75 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -47,6 +47,10 @@ household_support_summary, select_donor_households, ) +from policyengine_us_data.datasets.cps.long_term.prototype_synthetic_2100_support import ( + SyntheticCandidate, + build_role_donor_composites, +) class ExplodingCalibrator: @@ -339,6 +343,85 @@ def test_support_augmentation_appends_mixed_age_household(): assert synthetic_rows["tax_unit_id__2024"].nunique() == 3 +def test_role_donor_composites_build_structural_candidate_from_role_donors(): + import pandas as pd + + candidates = [ + SyntheticCandidate( + archetype="older_plus_prime_worker_family", + head_age=67, + spouse_age=42, + dependent_ages=(10,), + head_wages=0.0, + spouse_wages=100_000.0, + head_ss=40_000.0, + spouse_ss=0.0, + pension_income=0.0, + dividend_income=0.0, + ) + ] + actual_summary = pd.DataFrame( + [ + { + "tax_unit_id": 1, + "head_age": 70.0, + "spouse_age": None, + "adult_count": 1, + "dependent_count": 0, + "dependent_ages": (), + "head_payroll": 0.0, + "spouse_payroll": 0.0, + "head_ss": 40_000.0, + "spouse_ss": 0.0, + "payroll_total": 0.0, + "ss_total": 40_000.0, + "dividend_income": 2_000.0, + "pension_income": 8_000.0, + "support_count_weight": 1.0, + "person_weight_proxy": 1.0, + "archetype": "older_beneficiary_single", + }, + { + "tax_unit_id": 2, + "head_age": 41.0, + "spouse_age": 39.0, + "adult_count": 2, + "dependent_count": 1, + "dependent_ages": (10,), + "head_payroll": 60_000.0, + "spouse_payroll": 40_000.0, + "head_ss": 0.0, + "spouse_ss": 0.0, + "payroll_total": 100_000.0, + "ss_total": 0.0, + "dividend_income": 0.0, + "pension_income": 0.0, + "support_count_weight": 1.0, + "person_weight_proxy": 1.0, + "archetype": "prime_worker_family", + }, + ] + ) + + composite_candidates, prior_weights, report = build_role_donor_composites( + candidates, + np.array([1.0]), + actual_summary, + ss_scale=1.0, + earnings_scale=1.0, + top_n_targets=1, + older_donors_per_target=1, + worker_donors_per_target=1, + ) + + assert len(composite_candidates) == 1 + assert composite_candidates[0].archetype.endswith("_role_donor") + assert composite_candidates[0].spouse_wages == pytest.approx(100_000.0) + assert composite_candidates[0].head_ss == pytest.approx(40_000.0) + assert prior_weights.tolist() == pytest.approx([1.0]) + assert report["skipped_targets"] == [] + + def test_age_bin_helpers_preserve_population_totals(): bins = build_age_bins(n_ages=86, bucket_size=5) assert bins[0] == (0, 5) From 5643fddf7f56ac1b52754bc5221fed50325939a2 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 1 Apr 2026 11:54:50 -0400 Subject: [PATCH 24/31] Diagnose and fix support-augmentation translation --- .../datasets/cps/long_term/README.md | 13 +- .../cps/long_term/calibration_artifacts.py | 14 + ...agnose_support_augmentation_translation.py | 91 ++++ .../prototype_synthetic_2100_support.py | 446 +++++++++++++++++- .../cps/long_term/run_household_projection.py | 12 +- .../test_long_term_calibration_contract.py | 49 ++ 6 files changed, 599 insertions(+), 26 deletions(-) create mode 100644 policyengine_us_data/datasets/cps/long_term/diagnose_support_augmentation_translation.py diff --git a/policyengine_us_data/datasets/cps/long_term/README.md b/policyengine_us_data/datasets/cps/long_term/README.md index 8f1360d9e..60c21e660 100644 --- a/policyengine_us_data/datasets/cps/long_term/README.md +++ b/policyengine_us_data/datasets/cps/long_term/README.md @@ -56,6 +56,7 @@ python run_household_projection.py 2100 --profile ss profile/base dataset contract for the full artifact set. - Profiles validate achieved constraint errors before writing output. - Experimental donor-backed augmentation is stamped into each year sidecar and the directory manifest via `support_augmentation`. +- Donor-backed runs now also write a shared `support_augmentation_report.json` artifact with per-clone provenance so late-year translation failures can be inspected directly. **Estimated runtime:** ~2 minutes/year without `--save-h5`, ~3 minutes/year with `--save-h5` @@ -80,7 +81,7 @@ python run_household_projection.py 2100 --profile ss - Maps those synthetic targets back to nearest real 2024 donor tax units - Clones and perturbs the donor tax units to create a small augmented support without replacing the base CPS sample - Intended to test whether donor-backed synthetic support improves late-year microsim feasibility without resorting to fully free synthetic records -- Current status: integrated into the runner and fully auditable in metadata, but still diagnostic. The first `2100` end-to-end run did not materially improve the late-tail calibration frontier or support-concentration metrics. +- Current status: integrated into the runner and fully auditable in metadata, but still diagnostic. The first `2100` end-to-end run did not materially improve the late-tail calibration frontier or support-concentration metrics. After adding clone-provenance diagnostics and donor-specific inverse uprating, the realized clone households now match their intended `2100` age/SS/payroll targets closely; the remaining blocker is how much synthetic support is injected, not whether the actual-row translation works. **Role-Based Donor Composites** - Experimental structural extension of the donor-backed approach @@ -175,7 +176,15 @@ For each projection year (2025-2100): 5. **Benchmark TOB** - Compare modeled OASDI/HI TOB to the selected target source without forcing it into the weights 6. **Aggregate results** - Apply calibrated weights to calculate national totals -When donor-backed augmentation is enabled, step 1 uses the original 2024 CPS support and step 2 inserts a tagged late-year supplement derived from nearest real donors before the calibration loop begins. The underlying base dataset path remains unchanged in metadata; the augmentation details are recorded separately in `support_augmentation`. +When donor-backed augmentation is enabled, step 1 uses the original 2024 CPS support and step 2 inserts a tagged late-year supplement derived from nearest real donors before the calibration loop begins. The underlying base dataset path remains unchanged in metadata; the augmentation details are recorded separately in `support_augmentation`, and the full augmentation build report is written once per output directory to `support_augmentation_report.json`. + +To compare the intended clone targets with the realized output H5, run: + +```bash +uv run python policyengine_us_data/datasets/cps/long_term/diagnose_support_augmentation_translation.py \ + ./projected_datasets/2100.h5 \ + --year 2100 +``` **Key innovation:** Household-level calculations avoid person→household aggregation issues, maintaining consistency across all variables. diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py index 2676895b9..ece85be86 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py @@ -21,6 +21,7 @@ CONTRACT_VERSION = 1 MANIFEST_FILENAME = "calibration_manifest.json" +SUPPORT_AUGMENTATION_REPORT_FILENAME = "support_augmentation_report.json" def metadata_path_for(h5_path: str | Path) -> Path: @@ -124,6 +125,19 @@ def write_year_metadata( return metadata_path +def write_support_augmentation_report( + output_dir: str | Path, report: dict[str, Any] +) -> Path: + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + report_path = output_dir / SUPPORT_AUGMENTATION_REPORT_FILENAME + report_path.write_text( + json.dumps(json.loads(json.dumps(report)), indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return report_path + + def update_dataset_manifest( output_dir: str | Path, *, diff --git a/policyengine_us_data/datasets/cps/long_term/diagnose_support_augmentation_translation.py b/policyengine_us_data/datasets/cps/long_term/diagnose_support_augmentation_translation.py new file mode 100644 index 000000000..505a647d4 --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/diagnose_support_augmentation_translation.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +try: + from .prototype_synthetic_2100_support import summarize_realized_clone_translation +except ImportError: # pragma: no cover - script execution fallback + from prototype_synthetic_2100_support import summarize_realized_clone_translation + + +def _default_metadata_path(h5_path: Path) -> Path: + return h5_path.with_suffix(h5_path.suffix + ".metadata.json") + + +def main() -> int: + parser = argparse.ArgumentParser( + description=( + "Compare support-augmentation clone targets to the realized output H5." + ) + ) + parser.add_argument("h5_path", type=Path, help="Path to the realized year H5 file.") + parser.add_argument( + "--metadata", + type=Path, + default=None, + help="Path to the year metadata sidecar. Defaults to
.metadata.json.", + ) + parser.add_argument( + "--report", + type=Path, + default=None, + help=( + "Optional path to support_augmentation_report.json. Defaults to the " + "report_file named in metadata, or /support_augmentation_report.json." + ), + ) + parser.add_argument("--year", type=int, required=True, help="Output year to inspect.") + parser.add_argument( + "--age-bucket-size", + type=int, + default=5, + help="Age bucket size for the translation comparison summary.", + ) + parser.add_argument( + "--output", + type=Path, + default=None, + help="Optional JSON path to write the diagnostic summary.", + ) + args = parser.parse_args() + + metadata_path = args.metadata or _default_metadata_path(args.h5_path) + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + augmentation_metadata = metadata.get("support_augmentation") + if not augmentation_metadata: + raise ValueError(f"No support_augmentation metadata found in {metadata_path}") + + report_path = args.report + if report_path is None: + report_file = augmentation_metadata.get("report_file") + if report_file: + report_path = metadata_path.parent / report_file + else: + report_path = metadata_path.parent / "support_augmentation_report.json" + + augmentation_report = json.loads(report_path.read_text(encoding="utf-8")) + summary = summarize_realized_clone_translation( + str(args.h5_path), + period=args.year, + augmentation_report=augmentation_report, + age_bucket_size=args.age_bucket_size, + ) + payload = { + "h5_path": str(args.h5_path), + "metadata_path": str(metadata_path), + "report_path": str(report_path), + "year": int(args.year), + "age_bucket_size": int(args.age_bucket_size), + "summary": summary, + } + rendered = json.dumps(payload, indent=2, sort_keys=True) + "\n" + if args.output is not None: + args.output.write_text(rendered, encoding="utf-8") + print(rendered, end="") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py index 87fa6b8af..8450c5f9e 100644 --- a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py +++ b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py @@ -63,6 +63,8 @@ "employment_income_before_lsr", "self_employment_income_before_lsr", ) +PAYROLL_UPRATING_FACTOR_COLUMN = "__pe_payroll_uprating_factor" +SS_UPRATING_FACTOR_COLUMN = "__pe_ss_uprating_factor" @dataclass(frozen=True) @@ -417,35 +419,43 @@ def classify_archetype( return "prime_worker_single" if positive_payroll else "prime_other_single" -def build_actual_tax_unit_summary(base_dataset: str) -> pd.DataFrame: - sim = Microsimulation(dataset=base_dataset) +def build_tax_unit_summary( + dataset: str | Dataset, + *, + period: int, +) -> pd.DataFrame: + sim = Microsimulation(dataset=dataset) input_df = sim.to_input_dataframe() person_df = pd.DataFrame( { - "tax_unit_id": sim.calculate("person_tax_unit_id", period=BASE_YEAR).values, - "age": sim.calculate("age", period=BASE_YEAR).values, - "is_head": sim.calculate("is_tax_unit_head", period=BASE_YEAR).values, - "is_spouse": sim.calculate("is_tax_unit_spouse", period=BASE_YEAR).values, + "tax_unit_id": sim.calculate("person_tax_unit_id", period=period).values, + "household_id": sim.calculate("person_household_id", period=period).values, + "age": sim.calculate("age", period=period).values, + "is_head": sim.calculate("is_tax_unit_head", period=period).values, + "is_spouse": sim.calculate("is_tax_unit_spouse", period=period).values, "is_dependent": sim.calculate( - "is_tax_unit_dependent", period=BASE_YEAR + "is_tax_unit_dependent", period=period ).values, "social_security": sim.calculate( - "social_security", period=BASE_YEAR + "social_security", period=period ).values, "payroll": ( - sim.calculate("employment_income_before_lsr", period=BASE_YEAR).values + sim.calculate("employment_income_before_lsr", period=period).values + sim.calculate( - "self_employment_income_before_lsr", period=BASE_YEAR + "self_employment_income_before_lsr", period=period ).values ), "dividend_income": sim.calculate( - "qualified_dividend_income", period=BASE_YEAR + "qualified_dividend_income", period=period ).values, "pension_income": sim.calculate( - "taxable_pension_income", period=BASE_YEAR + "taxable_pension_income", period=period ).values, - "person_weight": input_df[f"person_weight__{BASE_YEAR}"].astype(float).values, + "person_weight": input_df[f"person_weight__{period}"].astype(float).values, + "household_weight": input_df[ + f"household_weight__{period}" + ].astype(float).values, } ) @@ -474,6 +484,7 @@ def build_actual_tax_unit_summary(base_dataset: str) -> pd.DataFrame: ) row = { "tax_unit_id": int(tax_unit_id), + "household_id": int(group["household_id"].iloc[0]), "head_age": head_age, "spouse_age": spouse_age, "adult_count": adult_count, @@ -489,6 +500,7 @@ def build_actual_tax_unit_summary(base_dataset: str) -> pd.DataFrame: "pension_income": float(group["pension_income"].sum()), "support_count_weight": 1.0, "person_weight_proxy": float(group["person_weight"].max()), + "household_weight_proxy": float(group["household_weight"].max()), } row["archetype"] = classify_archetype( head_age=row["head_age"], @@ -502,6 +514,59 @@ def build_actual_tax_unit_summary(base_dataset: str) -> pd.DataFrame: return pd.DataFrame(rows) +def build_actual_tax_unit_summary(base_dataset: str) -> pd.DataFrame: + return build_tax_unit_summary(base_dataset, period=BASE_YEAR) + + +def attach_person_uprating_factors( + input_df: pd.DataFrame, + sim: Microsimulation, + *, + base_year: int, + target_year: int, +) -> pd.DataFrame: + df = input_df.copy() + payroll_columns = [ + _period_column(component, base_year) + for component in PAYROLL_COMPONENTS + if _period_column(component, base_year) in df.columns + ] + ss_columns = [ + _period_column(component, base_year) + for component in SS_COMPONENTS + if _period_column(component, base_year) in df.columns + ] + base_payroll = ( + df[payroll_columns].astype(float).sum(axis=1).to_numpy() + if payroll_columns + else np.zeros(len(df), dtype=float) + ) + base_ss = ( + df[ss_columns].astype(float).sum(axis=1).to_numpy() + if ss_columns + else np.zeros(len(df), dtype=float) + ) + uprated_payroll = sum( + sim.calculate(component, period=target_year).values.astype(float) + for component in PAYROLL_COMPONENTS + ) + uprated_ss = sum( + sim.calculate(component, period=target_year).values.astype(float) + for component in SS_COMPONENTS + ) + df[PAYROLL_UPRATING_FACTOR_COLUMN] = np.where( + base_payroll > 0, + uprated_payroll / np.maximum(base_payroll, 1e-12), + np.nan, + ) + df[SS_UPRATING_FACTOR_COLUMN] = np.where( + base_ss > 0, + uprated_ss / np.maximum(base_ss, 1e-12), + np.nan, + ) + return df + + def load_base_aggregates(base_dataset: str) -> dict[str, float]: sim = Microsimulation(dataset=base_dataset) household_series = sim.calculate("household_id", period=BASE_YEAR, map_to="household") @@ -725,6 +790,227 @@ def age_bucket_vector(ages: list[int], age_bins: list[tuple[int, int]]) -> np.nd return vector +def _ages_from_summary_row(row: pd.Series) -> list[int]: + ages = [int(round(float(row["head_age"])))] + if pd.notna(row.get("spouse_age")): + ages.append(int(round(float(row["spouse_age"])))) + ages.extend(int(age) for age in row.get("dependent_ages", ())) + return ages + + +def _clone_report_record( + *, + clone_df: pd.DataFrame, + base_year: int, + target_candidate: SyntheticCandidate, + candidate_idx: int, + target_weight_share_pct: float, + clone_weight_scale: float, + combination_count: int, + older_donor_row: pd.Series | None, + worker_donor_row: pd.Series | None, +) -> dict[str, object]: + household_id_col = _period_column("household_id", base_year) + tax_unit_id_col = _period_column("tax_unit_id", base_year) + age_col = _period_column("age", base_year) + payroll_columns = [ + _period_column(component, base_year) + for component in PAYROLL_COMPONENTS + if _period_column(component, base_year) in clone_df.columns + ] + ss_columns = [ + _period_column(component, base_year) + for component in SS_COMPONENTS + if _period_column(component, base_year) in clone_df.columns + ] + ages = sorted(int(round(age)) for age in clone_df[age_col].astype(float).tolist()) + return { + "candidate_idx": int(candidate_idx), + "archetype": target_candidate.archetype, + "clone_household_id": int(clone_df[household_id_col].iloc[0]), + "clone_tax_unit_id": int(clone_df[tax_unit_id_col].iloc[0]), + "clone_person_count": int(len(clone_df)), + "clone_ages": ages, + "base_clone_payroll_total": float(clone_df[payroll_columns].sum().sum()) + if payroll_columns + else 0.0, + "base_clone_ss_total": float(clone_df[ss_columns].sum().sum()) + if ss_columns + else 0.0, + "target_weight_share_pct": float(target_weight_share_pct), + "per_clone_weight_share_pct": float( + target_weight_share_pct / max(combination_count, 1) + ), + "clone_weight_scale": float(clone_weight_scale), + "target_head_age": int(target_candidate.head_age), + "target_spouse_age": ( + int(target_candidate.spouse_age) + if target_candidate.spouse_age is not None + else None + ), + "target_dependent_ages": list(target_candidate.dependent_ages), + "target_head_wages": float(target_candidate.head_wages), + "target_spouse_wages": float(target_candidate.spouse_wages), + "target_head_ss": float(target_candidate.head_ss), + "target_spouse_ss": float(target_candidate.spouse_ss), + "target_payroll_total": float(target_candidate.payroll_total), + "target_ss_total": float(target_candidate.ss_total), + "older_donor_tax_unit_id": ( + int(older_donor_row["tax_unit_id"]) if older_donor_row is not None else None + ), + "worker_donor_tax_unit_id": ( + int(worker_donor_row["tax_unit_id"]) if worker_donor_row is not None else None + ), + "older_donor_distance": ( + float(older_donor_row["distance"]) + if older_donor_row is not None and "distance" in older_donor_row + else None + ), + "worker_donor_distance": ( + float(worker_donor_row["distance"]) + if worker_donor_row is not None and "distance" in worker_donor_row + else None + ), + } + + +def summarize_realized_clone_translation( + dataset: str | Dataset, + *, + period: int, + augmentation_report: dict[str, object], + age_bucket_size: int = 5, +) -> dict[str, object]: + clone_reports = augmentation_report.get("clone_household_reports", []) + if not clone_reports: + return { + "clone_household_count": 0, + "matched_clone_household_count": 0, + "unmatched_clone_household_count": 0, + "per_clone": [], + "by_archetype": [], + } + + realized_summary = build_tax_unit_summary(dataset, period=period) + realized_by_tax_unit = realized_summary.set_index("tax_unit_id", drop=False) + age_bins = build_age_bins(85, bucket_size=age_bucket_size) + per_clone: list[dict[str, object]] = [] + + for clone_report in clone_reports: + target_ages = [int(clone_report["target_head_age"])] + if clone_report.get("target_spouse_age") is not None: + target_ages.append(int(clone_report["target_spouse_age"])) + target_ages.extend(int(age) for age in clone_report["target_dependent_ages"]) + target_vector = age_bucket_vector(target_ages, age_bins) + + clone_tax_unit_id = int(clone_report["clone_tax_unit_id"]) + if clone_tax_unit_id not in realized_by_tax_unit.index: + per_clone.append( + { + **clone_report, + "matched": False, + "realized_archetype": None, + "realized_ages": None, + "realized_ss_total": None, + "realized_payroll_total": None, + "age_bucket_l1": None, + "ss_pct_error": None, + "payroll_pct_error": None, + } + ) + continue + + realized_row = realized_by_tax_unit.loc[clone_tax_unit_id] + realized_ages = _ages_from_summary_row(realized_row) + realized_vector = age_bucket_vector(realized_ages, age_bins) + realized_ss_total = float(realized_row["ss_total"]) + realized_payroll_total = float(realized_row["payroll_total"]) + target_ss_total = float(clone_report["target_ss_total"]) + target_payroll_total = float(clone_report["target_payroll_total"]) + per_clone.append( + { + **clone_report, + "matched": True, + "realized_archetype": realized_row["archetype"], + "realized_ages": realized_ages, + "realized_ss_total": realized_ss_total, + "realized_payroll_total": realized_payroll_total, + "realized_household_weight": float(realized_row["household_weight_proxy"]), + "age_bucket_l1": float(np.abs(realized_vector - target_vector).sum()), + "ss_pct_error": ( + 0.0 + if abs(target_ss_total) < 1e-9 + else (realized_ss_total - target_ss_total) / target_ss_total * 100 + ), + "payroll_pct_error": ( + 0.0 + if abs(target_payroll_total) < 1e-9 + else (realized_payroll_total - target_payroll_total) + / target_payroll_total + * 100 + ), + } + ) + + per_clone_df = pd.DataFrame(per_clone) + matched_df = per_clone_df[per_clone_df["matched"]].copy() + if matched_df.empty: + return { + "clone_household_count": int(len(per_clone_df)), + "matched_clone_household_count": 0, + "unmatched_clone_household_count": int(len(per_clone_df)), + "per_clone": per_clone, + "by_archetype": [], + } + + by_archetype = ( + matched_df.groupby("archetype", sort=False) + .agg( + clone_household_count=("clone_tax_unit_id", "count"), + avg_age_bucket_l1=("age_bucket_l1", "mean"), + avg_ss_pct_error=("ss_pct_error", "mean"), + avg_payroll_pct_error=("payroll_pct_error", "mean"), + ) + .reset_index() + ) + target_ss_total = float(matched_df["target_ss_total"].sum()) + target_payroll_total = float(matched_df["target_payroll_total"].sum()) + realized_ss_total = float(matched_df["realized_ss_total"].sum()) + realized_payroll_total = float(matched_df["realized_payroll_total"].sum()) + return { + "clone_household_count": int(len(per_clone_df)), + "matched_clone_household_count": int(len(matched_df)), + "unmatched_clone_household_count": int(len(per_clone_df) - len(matched_df)), + "target_ss_total": target_ss_total, + "realized_ss_total": realized_ss_total, + "aggregate_ss_pct_error": ( + 0.0 + if abs(target_ss_total) < 1e-9 + else (realized_ss_total - target_ss_total) / target_ss_total * 100 + ), + "target_payroll_total": target_payroll_total, + "realized_payroll_total": realized_payroll_total, + "aggregate_payroll_pct_error": ( + 0.0 + if abs(target_payroll_total) < 1e-9 + else (realized_payroll_total - target_payroll_total) + / target_payroll_total + * 100 + ), + "median_age_bucket_l1": float(matched_df["age_bucket_l1"].median()), + "median_ss_pct_error": float(matched_df["ss_pct_error"].median()), + "median_payroll_pct_error": float(matched_df["payroll_pct_error"].median()), + "top_ss_over": matched_df.sort_values("ss_pct_error", ascending=False) + .head(10) + .to_dict("records"), + "top_payroll_over": matched_df.sort_values("payroll_pct_error", ascending=False) + .head(10) + .to_dict("records"), + "by_archetype": by_archetype.to_dict("records"), + "per_clone": per_clone, + } + + def build_synthetic_constraint_problem( candidates: list[SyntheticCandidate], *, @@ -1310,6 +1596,22 @@ def _scale_person_components( return row +def _target_base_total_for_row( + row: pd.Series, + *, + target_total: float, + factor_column: str, + fallback_factor: float, +) -> float: + target_total = float(target_total) + if target_total <= 0: + return 0.0 + factor = row.get(factor_column, np.nan) + if pd.isna(factor) or float(factor) <= 0: + factor = fallback_factor + return target_total / max(float(factor), 1e-12) + + def _clone_tax_unit_rows_to_target( donor_rows: pd.DataFrame, *, @@ -1381,10 +1683,30 @@ def _clone_tax_unit_rows_to_target( ss_columns = tuple(_period_column(component, base_year) for component in SS_COMPONENTS) qbi_col = _period_column("w2_wages_from_qualified_business", base_year) - target_head_payroll = float(target_candidate.head_wages) / max(earnings_scale, 1e-12) - target_spouse_payroll = float(target_candidate.spouse_wages) / max(earnings_scale, 1e-12) - target_head_ss = float(target_candidate.head_ss) / max(ss_scale, 1e-12) - target_spouse_ss = float(target_candidate.spouse_ss) / max(ss_scale, 1e-12) + target_head_payroll = _target_base_total_for_row( + cloned.loc[head_idx], + target_total=float(target_candidate.head_wages), + factor_column=PAYROLL_UPRATING_FACTOR_COLUMN, + fallback_factor=earnings_scale, + ) + target_spouse_payroll = _target_base_total_for_row( + cloned.loc[spouse_idx] if spouse_idx is not None else pd.Series(dtype=float), + target_total=float(target_candidate.spouse_wages), + factor_column=PAYROLL_UPRATING_FACTOR_COLUMN, + fallback_factor=earnings_scale, + ) + target_head_ss = _target_base_total_for_row( + cloned.loc[head_idx], + target_total=float(target_candidate.head_ss), + factor_column=SS_UPRATING_FACTOR_COLUMN, + fallback_factor=ss_scale, + ) + target_spouse_ss = _target_base_total_for_row( + cloned.loc[spouse_idx] if spouse_idx is not None else pd.Series(dtype=float), + target_total=float(target_candidate.spouse_ss), + factor_column=SS_UPRATING_FACTOR_COLUMN, + fallback_factor=ss_scale, + ) cloned.loc[head_idx] = _scale_person_components( cloned.loc[head_idx].copy(), @@ -1461,6 +1783,36 @@ def _dependent_rows(df: pd.DataFrame | None) -> pd.DataFrame: worker_adults = _adult_rows(worker_donor_rows) worker_dependents = _dependent_rows(worker_donor_rows) + worker_payroll_factor = ( + float( + np.nanmedian( + worker_adults[PAYROLL_UPRATING_FACTOR_COLUMN].astype(float).to_numpy() + ) + ) + if not worker_adults.empty + and PAYROLL_UPRATING_FACTOR_COLUMN in worker_adults.columns + else np.nan + ) + older_ss_factor = ( + float( + np.nanmedian( + older_adults[SS_UPRATING_FACTOR_COLUMN].astype(float).to_numpy() + ) + ) + if not older_adults.empty and SS_UPRATING_FACTOR_COLUMN in older_adults.columns + else np.nan + ) + payroll_reference_factor = ( + worker_payroll_factor + if np.isfinite(worker_payroll_factor) and worker_payroll_factor > 0 + else earnings_scale + ) + ss_reference_factor = ( + older_ss_factor + if np.isfinite(older_ss_factor) and older_ss_factor > 0 + else ss_scale + ) + selected_rows: list[pd.Series] = [] head_target_older = target_candidate.head_ss > 0 or target_candidate.head_age >= 65 head_source_rows = older_adults if head_target_older and not older_adults.empty else worker_adults @@ -1548,10 +1900,19 @@ def _dependent_rows(df: pd.DataFrame | None) -> pd.DataFrame: ss_columns = tuple(_period_column(component, base_year) for component in SS_COMPONENTS) qbi_col = _period_column("w2_wages_from_qualified_business", base_year) - target_head_payroll = float(target_candidate.head_wages) / max(earnings_scale, 1e-12) - target_spouse_payroll = float(target_candidate.spouse_wages) / max(earnings_scale, 1e-12) - target_head_ss = float(target_candidate.head_ss) / max(ss_scale, 1e-12) - target_spouse_ss = float(target_candidate.spouse_ss) / max(ss_scale, 1e-12) + target_head_payroll = float(target_candidate.head_wages) / max( + payroll_reference_factor, + 1e-12, + ) + target_spouse_payroll = float(target_candidate.spouse_wages) / max( + payroll_reference_factor, + 1e-12, + ) + target_head_ss = float(target_candidate.head_ss) / max(ss_reference_factor, 1e-12) + target_spouse_ss = float(target_candidate.spouse_ss) / max( + ss_reference_factor, + 1e-12, + ) cloned.loc[head_idx] = _scale_person_components( cloned.loc[head_idx].copy(), @@ -1607,7 +1968,12 @@ def build_donor_backed_augmented_input_dataframe( clone_weight_scale: float = 0.1, ) -> tuple[pd.DataFrame, dict[str, object]]: sim = Microsimulation(dataset=base_dataset) - input_df = sim.to_input_dataframe() + input_df = attach_person_uprating_factors( + sim.to_input_dataframe(), + sim, + base_year=base_year, + target_year=target_year, + ) actual_summary = build_actual_tax_unit_summary(base_dataset) base_aggregates = load_base_aggregates(base_dataset) ss_scale = load_ssa_benefit_projections(target_year) / max( @@ -1641,6 +2007,7 @@ def build_donor_backed_augmented_input_dataframe( id_counters["person"] = _next_entity_id(input_df[_period_column(PERSON_ID_COLUMN, base_year)]) clone_frames = [] + clone_household_reports = [] target_reports = [] skipped_targets = [] @@ -1695,6 +2062,12 @@ def build_donor_backed_augmented_input_dataframe( if clone_frames else input_df.copy() ) + helper_columns = [PAYROLL_UPRATING_FACTOR_COLUMN, SS_UPRATING_FACTOR_COLUMN] + augmented_df.drop( + columns=[column for column in helper_columns if column in augmented_df.columns], + inplace=True, + errors="ignore", + ) report = { "base_dataset": base_dataset, "base_year": int(base_year), @@ -1727,7 +2100,12 @@ def build_role_composite_augmented_input_dataframe( clone_weight_scale: float = 0.1, ) -> tuple[pd.DataFrame, dict[str, object]]: sim = Microsimulation(dataset=base_dataset) - input_df = sim.to_input_dataframe() + input_df = attach_person_uprating_factors( + sim.to_input_dataframe(), + sim, + base_year=base_year, + target_year=target_year, + ) actual_summary = build_actual_tax_unit_summary(base_dataset) base_aggregates = load_base_aggregates(base_dataset) ss_scale = load_ssa_benefit_projections(target_year) / max( @@ -1763,6 +2141,7 @@ def build_role_composite_augmented_input_dataframe( ) clone_frames = [] + clone_household_reports = [] target_reports = [] skipped_targets = [] @@ -1834,6 +2213,19 @@ def build_role_composite_augmented_input_dataframe( if clone_df is None: continue clone_frames.append(clone_df) + clone_household_reports.append( + _clone_report_record( + clone_df=clone_df, + base_year=base_year, + target_candidate=target_candidate, + candidate_idx=int(target_row["candidate_idx"]), + target_weight_share_pct=float(target_row["weight_share_pct"]), + clone_weight_scale=clone_weight_scale, + combination_count=combination_count, + older_donor_row=older_row if target_candidate.ss_total > 0 else None, + worker_donor_row=worker_row if target_candidate.payroll_total > 0 else None, + ) + ) successful_clone_count += 1 target_reports.append( { @@ -1851,6 +2243,12 @@ def build_role_composite_augmented_input_dataframe( if clone_frames else input_df.copy() ) + helper_columns = [PAYROLL_UPRATING_FACTOR_COLUMN, SS_UPRATING_FACTOR_COLUMN] + augmented_df.drop( + columns=[column for column in helper_columns if column in augmented_df.columns], + inplace=True, + errors="ignore", + ) report = { "base_dataset": base_dataset, "base_year": int(base_year), @@ -1870,6 +2268,8 @@ def build_role_composite_augmented_input_dataframe( ), "base_person_count": int(len(input_df)), "augmented_person_count": int(len(augmented_df)), + "clone_household_count": int(len(clone_household_reports)), + "clone_household_reports": clone_household_reports, "target_reports": target_reports, "skipped_targets": skipped_targets, } diff --git a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py index a14614b8f..ba54bc839 100644 --- a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py +++ b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py @@ -46,7 +46,11 @@ set_long_term_target_source, ) from calibration import build_calibration_audit, calibrate_weights -from calibration_artifacts import update_dataset_manifest, write_year_metadata +from calibration_artifacts import ( + update_dataset_manifest, + write_support_augmentation_report, + write_year_metadata, +) from calibration_profiles import ( approximate_window_for_year, build_profile_from_flags, @@ -561,6 +565,10 @@ def create_h6_reform(): max_worker_distance=SUPPORT_AUGMENTATION_MAX_DISTANCE, clone_weight_scale=SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE, ) + support_augmentation_report_path = write_support_augmentation_report( + OUTPUT_DIR, + augmentation_report, + ) SUPPORT_AUGMENTATION_METADATA = { "name": SUPPORT_AUGMENTATION_PROFILE, "activation_start_year": SUPPORT_AUGMENTATION_START_YEAR, @@ -569,6 +577,7 @@ def create_h6_reform(): "donors_per_target": SUPPORT_AUGMENTATION_DONORS_PER_TARGET, "max_distance_for_clone": SUPPORT_AUGMENTATION_MAX_DISTANCE, "clone_weight_scale": SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE, + "report_file": support_augmentation_report_path.name, "report_summary": { "base_household_count": augmentation_report["base_household_count"], "augmented_household_count": augmentation_report[ @@ -578,6 +587,7 @@ def create_h6_reform(): "augmented_person_count": augmentation_report[ "augmented_person_count" ], + "clone_household_count": augmentation_report.get("clone_household_count", 0), "successful_target_count": sum( report["successful_clone_count"] > 0 for report in augmentation_report["target_reports"] diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index 50e278a75..3236e9be1 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -3,6 +3,7 @@ import json import numpy as np import pytest +from policyengine_core.data.dataset import Dataset from policyengine_us_data.datasets.cps.long_term import ( calibration as calibration_module, @@ -18,6 +19,7 @@ normalize_metadata, rebuild_dataset_manifest, update_dataset_manifest, + write_support_augmentation_report, write_year_metadata, ) from policyengine_us_data.datasets.cps.long_term.calibration_profiles import ( @@ -50,6 +52,7 @@ from policyengine_us_data.datasets.cps.long_term.prototype_synthetic_2100_support import ( SyntheticCandidate, build_role_donor_composites, + summarize_realized_clone_translation, ) @@ -1053,6 +1056,7 @@ def test_manifest_persists_support_augmentation_metadata(tmp_path): "name": "donor-backed-synthetic-v1", "activation_start_year": 2075, "target_year": 2100, + "report_file": "support_augmentation_report.json", "report_summary": { "base_household_count": 41314, "augmented_household_count": 41326, @@ -1083,9 +1087,54 @@ def test_manifest_persists_support_augmentation_metadata(tmp_path): metadata = json.loads(metadata_path.read_text(encoding="utf-8")) manifest = json.loads(manifest_path.read_text(encoding="utf-8")) assert metadata["support_augmentation"]["name"] == "donor-backed-synthetic-v1" + assert metadata["support_augmentation"]["report_file"] == "support_augmentation_report.json" assert ( manifest["support_augmentation"]["report_summary"][ "augmented_household_count" ] == 41326 ) + + +def test_write_support_augmentation_report(tmp_path): + report = { + "name": "donor-backed-composite-v1", + "clone_household_count": 2, + "clone_household_reports": [{"clone_household_id": 1001}], + } + report_path = write_support_augmentation_report(tmp_path, report) + assert report_path == tmp_path / "support_augmentation_report.json" + loaded = json.loads(report_path.read_text(encoding="utf-8")) + assert loaded["clone_household_count"] == 2 + assert loaded["clone_household_reports"][0]["clone_household_id"] == 1001 + + +def test_summarize_realized_clone_translation_matches_toy_clone(): + import pandas as pd + + dataset = Dataset.from_dataframe(pd.DataFrame(_toy_support_dataframe()), 2024) + augmentation_report = { + "clone_household_reports": [ + { + "candidate_idx": 0, + "archetype": "older_worker_couple", + "clone_household_id": 1, + "clone_tax_unit_id": 101, + "target_head_age": 70, + "target_spouse_age": 68, + "target_dependent_ages": [], + "target_payroll_total": 5_000.0, + "target_ss_total": 20_000.0, + } + ] + } + summary = summarize_realized_clone_translation( + dataset, + period=2024, + augmentation_report=augmentation_report, + age_bucket_size=5, + ) + assert summary["matched_clone_household_count"] == 1 + assert summary["aggregate_ss_pct_error"] == pytest.approx(0.0) + assert summary["aggregate_payroll_pct_error"] == pytest.approx(0.0) + assert summary["per_clone"][0]["realized_ages"] == [70, 68] From 0b7dfce47505f0af6f5b1db272a63bbecee63aed Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 2 Apr 2026 07:17:17 -0400 Subject: [PATCH 25/31] Add long-run TOB comparison note --- .../cps/long_term/TOB_ALIGNMENT_NOTE.md | 100 ++++++++++++++++ .../cps/long_term/compare_tob_shares.py | 111 ++++++++++++++++++ 2 files changed, 211 insertions(+) create mode 100644 policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md create mode 100644 policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py diff --git a/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md b/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md new file mode 100644 index 000000000..c487b92f8 --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md @@ -0,0 +1,100 @@ +# Long-Run TOB Alignment Note + +This note records the current apples-to-apples comparison point for long-run +taxation of benefits (TOB). + +## Why this note exists + +Our earlier comparison mixed two different objects: + +- `OASDI-only` income from taxation of benefits as a share of total OASDI + benefits +- combined `OASDI + HI` credited taxes as a share of total OASDI benefits + +The public Urban / DYNASIM and SSA Trustees discussion is about the former, +not the latter. + +## Primary-source baseline + +- SSA 2025 Trustees Report, [V.C.7](https://www.ssa.gov/oact/tr/2025/V_C_prog.html) + says the benefit-tax thresholds are "constant in the future" and, after the + tenth projection year, ordinary income-tax brackets "rise with average + wages." +- Urban's 2024 DYNASIM appendix says it continues current indexing of income + tax parameters indefinitely and keeps the Social Security benefit-tax + thresholds at current nominal levels throughout the projection period. Source: + [Urban 2024 appendix](https://www.urban.org/sites/default/files/2024-10/Does-the-2023-Social-Security-Expansion-Act-Improve-Equity-in-Key-Outcomes.pdf). +- SSA's published OASDI long-run target series in + [Table IV.B2](https://www.ssa.gov/oact/tr/2025/lr4b2.html) and our local + `trustees_2025_current_law.csv` target package imply OASDI-only TOB shares + of about `6.0%` to `6.1%` of OASDI benefits in the late horizon. + +## Current branch-local comparison + +The table below uses one-year probe outputs produced on `2026-04-02` with: + +- `policyengine-us-data` branch `codex/us-data-calibration-contract` +- `policyengine-us` branch `codex/extend-ss-cap-2100` or equivalent fix from + [PR #7912](https://github.com/PolicyEngine/policyengine-us/pull/7912) +- profile `ss-payroll-tob` +- target source `trustees_2025_current_law` +- donor-composite support augmentation enabled + +| Year | OASDI actual | OASDI target | OASDI gap | Combined actual | Combined target | Combined gap | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| 2075 | 9.43% | 6.01% | +3.42 pp | 12.58% | 11.09% | +1.50 pp | +| 2090 | 10.52% | 6.08% | +4.44 pp | 16.66% | 11.20% | +5.46 pp | +| 2100 | 11.16% | 6.10% | +5.06 pp | 19.35% | 11.22% | +8.14 pp | + +Interpretation: + +- The `OASDI-only` comparison is the right one for evaluating alignment to the + public Urban / SSA discussion. +- On that comparable metric, the current corrected baseline is above both the + Trustees-style target path (`~6.1%`) and Urban's public DYNASIM endpoint + (`8.5%` in `2095`). +- The larger combined `OASDI + HI` shares are still useful internal diagnostics, + but they should not be compared directly to the public `5.6%` / `8.5%` + figures. + +## DYNASIM public benchmark + +Urban's 2024 appendix says DYNASIM's revenue from taxing Social Security +benefits rises from `5 percent` in `2027` to `8.5 percent` in `2095`, while +the Social Security actuaries' corresponding share rises from `5 percent` to +`5.6 percent` over the same period. + +Important caveats: + +- We have not found a public annual DYNASIM TOB series, only these endpoint + shares. +- The Urban paper is tied to a 2023-vintage Trustees baseline. Comparing the + 2023 and 2025 Trustees reports, the long-run CPI and average-wage growth + assumptions are effectively unchanged, so that vintage difference does not + explain most of our remaining gap. + +## Reproducing the table + +The comparison script is: + +- [compare_tob_shares.py](./compare_tob_shares.py) + +Example: + +```bash +uv run python policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py \ + /path/to/2075-output-dir \ + /path/to/2090-output-dir \ + /path/to/2100-output-dir +``` + +The script expects metadata sidecars containing: + +- `calibration_audit.constraints.ss_total` +- `calibration_audit.benchmarks.oasdi_tob` +- `calibration_audit.benchmarks.hi_tob` + +To regenerate the underlying sidecars from scratch, first ensure that +`policyengine-us` includes the Social Security wage-base extension from +[PR #7912](https://github.com/PolicyEngine/policyengine-us/pull/7912) or later. +Without that fix, late-year taxable-payroll calibration is materially wrong. diff --git a/policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py b/policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py new file mode 100644 index 000000000..e3a7e6a4f --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +import argparse +import json +from pathlib import Path + + +def _resolve_metadata_paths(inputs: list[str]) -> list[Path]: + paths: list[Path] = [] + for raw in inputs: + path = Path(raw).expanduser() + if path.is_dir(): + for candidate in sorted(path.glob("*.metadata.json")): + paths.append(candidate) + continue + if path.is_file(): + paths.append(path) + continue + raise FileNotFoundError(f"Metadata path not found: {path}") + if not paths: + raise ValueError("No metadata files found.") + return paths + + +def _load_record(path: Path) -> dict: + metadata = json.loads(path.read_text(encoding="utf-8")) + audit = metadata["calibration_audit"] + constraints = audit["constraints"] + benchmarks = audit["benchmarks"] + + ss_actual = float(constraints["ss_total"]["achieved"]) + ss_target = float(constraints["ss_total"]["target"]) + oasdi_actual = float(benchmarks["oasdi_tob"]["achieved"]) + oasdi_target = float(benchmarks["oasdi_tob"]["target"]) + hi_actual = float(benchmarks["hi_tob"]["achieved"]) + hi_target = float(benchmarks["hi_tob"]["target"]) + + combined_actual = oasdi_actual + hi_actual + combined_target = oasdi_target + hi_target + + return { + "year": int(metadata["year"]), + "source_path": str(path), + "oasdi_actual_share_pct": 100 * oasdi_actual / ss_actual, + "oasdi_target_share_pct": 100 * oasdi_target / ss_target, + "oasdi_gap_pct_pt": 100 * oasdi_actual / ss_actual + - 100 * oasdi_target / ss_target, + "combined_actual_share_pct": 100 * combined_actual / ss_actual, + "combined_target_share_pct": 100 * combined_target / ss_target, + "combined_gap_pct_pt": 100 * combined_actual / ss_actual + - 100 * combined_target / ss_target, + } + + +def _format_markdown(records: list[dict]) -> str: + header = ( + "| Year | OASDI actual | OASDI target | OASDI gap | Combined actual | " + "Combined target | Combined gap |\n" + "| --- | ---: | ---: | ---: | ---: | ---: | ---: |" + ) + rows = [ + ( + f"| {record['year']} | " + f"{record['oasdi_actual_share_pct']:.2f}% | " + f"{record['oasdi_target_share_pct']:.2f}% | " + f"{record['oasdi_gap_pct_pt']:+.2f} pp | " + f"{record['combined_actual_share_pct']:.2f}% | " + f"{record['combined_target_share_pct']:.2f}% | " + f"{record['combined_gap_pct_pt']:+.2f} pp |" + ) + for record in records + ] + return "\n".join([header, *rows]) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Summarize OASDI-only and combined TOB shares from long-run " + "metadata sidecars." + ), + ) + parser.add_argument( + "paths", + nargs="+", + help=( + "Metadata files or directories containing *.metadata.json sidecars." + ), + ) + parser.add_argument( + "--format", + choices=("markdown", "json"), + default="markdown", + help="Output format.", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + paths = _resolve_metadata_paths(args.paths) + records = sorted((_load_record(path) for path in paths), key=lambda x: x["year"]) + if args.format == "json": + print(json.dumps(records, indent=2, sort_keys=True)) + else: + print(_format_markdown(records)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 40f8e450f39eeecdf55c2a68877293e1c8486a41 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 2 Apr 2026 07:50:13 -0400 Subject: [PATCH 26/31] Add Trustees bracket-indexing TOB benchmark --- .../cps/long_term/TOB_ALIGNMENT_NOTE.md | 39 +++ .../benchmark_trustees_bracket_indexing.py | 226 ++++++++++++++++++ 2 files changed, 265 insertions(+) create mode 100644 policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py diff --git a/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md b/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md index c487b92f8..f3bd2ea6c 100644 --- a/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md +++ b/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md @@ -57,6 +57,38 @@ Interpretation: but they should not be compared directly to the public `5.6%` / `8.5%` figures. +## Trustees-style bracket-growth sensitivity + +We also ran a narrow tax-side sensitivity that keeps the calibrated household +weights fixed and changes only one assumption in the tax model: + +- after `2034`, ordinary federal income-tax bracket thresholds are uprated with + `NAWI` instead of `C-CPI-U` +- the Social Security benefit-tax thresholds remain fixed + +This is intended as a best-public approximation to the Trustees statement that +ordinary income-tax brackets "rise with average wages" after the tenth +projection year. + +| Year | Baseline OASDI | Wage-indexed-brackets OASDI | Trustees target | Remaining gap | +| --- | ---: | ---: | ---: | ---: | +| 2075 | 9.43% | 8.41% | 6.01% | +2.40 pp | +| 2090 | 10.52% | 7.85% | 6.08% | +1.76 pp | +| 2100 | 11.16% | 9.46% | 6.10% | +3.36 pp | + +Interpretation: + +- This tax-side assumption moves the modeled OASDI-only TOB share materially in + the right direction. +- It explains a substantial share of the excess over the Trustees target, + especially around `2090`. +- It does not explain the whole gap. Even with wage-indexed ordinary brackets, + the long-run `2100` OASDI-only share remains well above the Trustees-style + `~6.1%` path. +- Relative to Urban's public DYNASIM endpoint of `8.5%` in `2095`, the + wage-indexed-brackets sensitivity lands in the same rough range by `2090`, + but is still above that public number by `2100`. + ## DYNASIM public benchmark Urban's 2024 appendix says DYNASIM's revenue from taxing Social Security @@ -78,6 +110,7 @@ Important caveats: The comparison script is: - [compare_tob_shares.py](./compare_tob_shares.py) +- [benchmark_trustees_bracket_indexing.py](./benchmark_trustees_bracket_indexing.py) Example: @@ -86,6 +119,12 @@ uv run python policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py /path/to/2075-output-dir \ /path/to/2090-output-dir \ /path/to/2100-output-dir + +uv run python policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py \ + /path/to/2075-output-dir \ + /path/to/2090-output-dir \ + /path/to/2100-output-dir \ + --policyengine-us-path /path/to/local/policyengine-us ``` The script expects metadata sidecars containing: diff --git a/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py b/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py new file mode 100644 index 000000000..ed5f3aaab --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py @@ -0,0 +1,226 @@ +from __future__ import annotations + +import argparse +import json +import math +import sys +from pathlib import Path + + +FILING_STATUSES = [ + "SINGLE", + "JOINT", + "SEPARATE", + "HEAD_OF_HOUSEHOLD", + "SURVIVING_SPOUSE", +] + + +def round_down(amount: float, interval: float) -> float: + return math.floor(amount / interval) * interval + + +def create_wage_indexed_brackets_reform( + start_year: int = 2035, + end_year: int = 2100, +): + from policyengine_us.model_api import Reform + + def modify_parameters(parameters): + nawi = parameters.gov.ssa.nawi + thresholds = parameters.gov.irs.income.bracket.thresholds + + for bracket in map(str, range(1, 7)): + bracket_node = thresholds.get_child(bracket) + for filing_status in FILING_STATUSES: + parameter = bracket_node.get_child(filing_status) + interval = float( + parameter.metadata["uprating"]["rounding"]["interval"] + ) + + for year in range(start_year, end_year + 1): + previous_value = float(parameter(f"{year - 1}-01-01")) + wage_growth = float(nawi(f"{year - 1}-01-01")) / float( + nawi(f"{year - 2}-01-01") + ) + updated_value = round_down(previous_value * wage_growth, interval) + parameter.update( + period=f"year:{year}-01-01:1", + value=updated_value, + ) + return parameters + + class reform(Reform): + def apply(self): + self.modify_parameters(modify_parameters) + + return reform + + +def _coerce_h5_path(raw: str) -> Path: + path = Path(raw).expanduser() + if path.is_dir(): + matches = sorted(path.glob("*.h5")) + if len(matches) != 1: + raise ValueError( + f"Expected exactly one .h5 file in {path}, found {len(matches)}" + ) + return matches[0] + if path.suffix == ".metadata.json": + return path.with_suffix("").with_suffix(".h5") + return path + + +def _load_metadata(h5_path: Path) -> dict | None: + metadata_path = h5_path.with_suffix(".h5.metadata.json") + if not metadata_path.exists(): + return None + return json.loads(metadata_path.read_text(encoding="utf-8")) + + +def _baseline_record(h5_path: Path, metadata: dict | None) -> dict: + if metadata is None: + return { + "year": int(h5_path.stem), + "baseline_oasdi_share_pct": None, + "baseline_combined_share_pct": None, + "target_oasdi_share_pct": None, + } + + audit = metadata["calibration_audit"] + ss_total = float(audit["constraints"]["ss_total"]["achieved"]) + ss_target = float(audit["constraints"]["ss_total"]["target"]) + oasdi_actual = float(audit["benchmarks"]["oasdi_tob"]["achieved"]) + hi_actual = float(audit["benchmarks"]["hi_tob"]["achieved"]) + oasdi_target = float(audit["benchmarks"]["oasdi_tob"]["target"]) + + return { + "year": int(metadata["year"]), + "baseline_oasdi_share_pct": 100 * oasdi_actual / ss_total, + "baseline_combined_share_pct": 100 * (oasdi_actual + hi_actual) / ss_total, + "target_oasdi_share_pct": 100 * oasdi_target / ss_target, + } + + +def _compute_reformed_shares( + h5_path: Path, + start_year: int, + end_year: int, +) -> dict: + from policyengine_us import Microsimulation + + reform = create_wage_indexed_brackets_reform( + start_year=start_year, + end_year=end_year, + ) + sim = Microsimulation(dataset=str(h5_path), reform=reform) + ss_total = float(sim.calculate("social_security").sum()) + oasdi_tob = float(sim.calculate("tob_revenue_oasdi").sum()) + hi_tob = float(sim.calculate("tob_revenue_medicare_hi").sum()) + return { + "reformed_oasdi_share_pct": 100 * oasdi_tob / ss_total, + "reformed_combined_share_pct": 100 * (oasdi_tob + hi_tob) / ss_total, + } + + +def _format_markdown(records: list[dict]) -> str: + header = ( + "| Year | Trustees OASDI target | Baseline OASDI | Wage-indexed OASDI | " + "OASDI delta | Baseline combined | Wage-indexed combined |\n" + "| --- | ---: | ---: | ---: | ---: | ---: | ---: |" + ) + rows = [] + for record in records: + baseline_oasdi = record["baseline_oasdi_share_pct"] + reformed_oasdi = record["reformed_oasdi_share_pct"] + baseline_combined = record["baseline_combined_share_pct"] + reformed_combined = record["reformed_combined_share_pct"] + rows.append( + "| {year} | {target:.2f}% | {base_o:.2f}% | {reform_o:.2f}% | {delta:+.2f} pp | " + "{base_c:.2f}% | {reform_c:.2f}% |".format( + year=record["year"], + target=record["target_oasdi_share_pct"], + base_o=baseline_oasdi, + reform_o=reformed_oasdi, + delta=reformed_oasdi - baseline_oasdi, + base_c=baseline_combined, + reform_c=reformed_combined, + ) + ) + return "\n".join([header, *rows]) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Benchmark calibrated long-run H5s under a wage-indexed ordinary " + "income-tax bracket sensitivity." + ), + ) + parser.add_argument( + "paths", + nargs="+", + help="H5 files, metadata files, or directories containing a single H5.", + ) + parser.add_argument( + "--policyengine-us-path", + help=( + "Optional local policyengine-us checkout to prepend to sys.path. " + "Use this when the required tax-side fix is not yet released." + ), + ) + parser.add_argument( + "--start-year", + type=int, + default=2035, + help="First year to switch ordinary federal brackets from CPI to wages.", + ) + parser.add_argument( + "--end-year", + type=int, + default=2100, + help="Last year to extend the wage-indexed bracket sensitivity through.", + ) + parser.add_argument( + "--format", + choices=("markdown", "json"), + default="markdown", + help="Output format.", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if args.policyengine_us_path: + sys.path.insert(0, str(Path(args.policyengine_us_path).expanduser())) + + h5_paths = [_coerce_h5_path(path) for path in args.paths] + records = [] + + for h5_path in h5_paths: + metadata = _load_metadata(h5_path) + baseline = _baseline_record(h5_path, metadata) + print( + f"[{baseline['year']}] benchmarking wage-indexed federal brackets on {h5_path}", + file=sys.stderr, + flush=True, + ) + reformed = _compute_reformed_shares( + h5_path, + start_year=args.start_year, + end_year=args.end_year, + ) + baseline.update(reformed) + records.append(baseline) + + records.sort(key=lambda record: record["year"]) + if args.format == "json": + print(json.dumps(records, indent=2, sort_keys=True)) + else: + print(_format_markdown(records)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From f35107d06349cac26436d0c6d2fada32be476fa8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 2 Apr 2026 08:16:32 -0400 Subject: [PATCH 27/31] Extend TOB tax-side benchmark scenarios --- .../cps/long_term/TOB_ALIGNMENT_NOTE.md | 32 ++++ .../benchmark_trustees_bracket_indexing.py | 170 ++++++++++++++++-- 2 files changed, 190 insertions(+), 12 deletions(-) diff --git a/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md b/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md index f3bd2ea6c..06c01f33a 100644 --- a/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md +++ b/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md @@ -89,6 +89,38 @@ Interpretation: wage-indexed-brackets sensitivity lands in the same rough range by `2090`, but is still above that public number by `2100`. +## Broader core-threshold sensitivity + +We also ran a broader but still targeted tax-side sensitivity that switches a +core set of federal thresholds from `C-CPI-U` to `NAWI` after `2034`: + +- ordinary income-tax brackets +- standard deduction +- aged/blind additional standard deduction +- capital-gains rate thresholds +- AMT bracket threshold and exemption thresholds + +This is broader than the minimum public Trustees approximation, but still +narrower than switching the entire `gov.irs.uprating` family to wages. + +| Year | Baseline OASDI | Core-threshold OASDI | Trustees target | Remaining gap | +| --- | ---: | ---: | ---: | ---: | +| 2075 | 9.43% | 7.65% | 6.01% | +1.64 pp | +| 2090 | 10.52% | 7.31% | 6.08% | +1.22 pp | +| 2100 | 11.16% | 8.15% | 6.10% | +2.05 pp | + +Interpretation: + +- The broader threshold bundle explains more of the TOB gap than brackets + alone. +- The additional movement is meaningful, especially in `2100`, where the + OASDI-only share falls from `9.46%` under brackets-only to `8.15%` under the + broader core-threshold sensitivity. +- Even this broader sensitivity still does not fully reconcile the modeled TOB + path to the Trustees target, so some remaining gap likely reflects + beneficiary income mix, filing composition, or other Treasury-ratio modeling + differences. + ## DYNASIM public benchmark Urban's 2024 appendix says DYNASIM's revenue from taxing Social Security diff --git a/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py b/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py index ed5f3aaab..95a11501c 100644 --- a/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py +++ b/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py @@ -20,6 +20,53 @@ def round_down(amount: float, interval: float) -> float: return math.floor(amount / interval) * interval +def round_amount(amount: float, rounding: dict | None) -> float: + if not rounding: + return amount + + interval = float(rounding["interval"]) + rounding_type = rounding["type"] + + if rounding_type == "downwards": + return math.floor(amount / interval) * interval + if rounding_type == "nearest": + return math.floor(amount / interval + 0.5) * interval + + raise ValueError(f"Unsupported rounding type: {rounding_type}") + + +def _iter_updatable_parameters(root) -> list: + candidates = [root] + if hasattr(root, "get_descendants"): + candidates.extend(root.get_descendants()) + + result = [] + for candidate in candidates: + if candidate.__class__.__name__ != "Parameter": + continue + uprating = getattr(candidate, "metadata", {}).get("uprating") + if uprating is not None: + result.append(candidate) + return result + + +def _apply_wage_growth_to_parameter(parameter, nawi, start_year: int, end_year: int): + metadata = getattr(parameter, "metadata", {}) + uprating = metadata.get("uprating") + rounding = uprating.get("rounding") if isinstance(uprating, dict) else None + + for year in range(start_year, end_year + 1): + previous_value = float(parameter(f"{year - 1}-01-01")) + wage_growth = float(nawi(f"{year - 1}-01-01")) / float( + nawi(f"{year - 2}-01-01") + ) + updated_value = round_amount(previous_value * wage_growth, rounding) + parameter.update( + period=f"year:{year}-01-01:1", + value=updated_value, + ) + + def create_wage_indexed_brackets_reform( start_year: int = 2035, end_year: int = 2100, @@ -57,6 +104,74 @@ def apply(self): return reform +def create_wage_indexed_core_thresholds_reform( + start_year: int = 2035, + end_year: int = 2100, +): + from policyengine_us.model_api import Reform + + def modify_parameters(parameters): + nawi = parameters.gov.ssa.nawi + roots = [ + parameters.gov.irs.income.bracket.thresholds, + parameters.gov.irs.deductions.standard.amount, + parameters.gov.irs.deductions.standard.aged_or_blind.amount, + parameters.gov.irs.capital_gains.thresholds, + parameters.gov.irs.income.amt.brackets, + parameters.gov.irs.income.amt.exemption.amount, + parameters.gov.irs.income.amt.exemption.phase_out.start, + parameters.gov.irs.income.amt.exemption.separate_limit, + ] + + seen = set() + for root in roots: + for parameter in _iter_updatable_parameters(root): + if parameter.name in seen: + continue + seen.add(parameter.name) + _apply_wage_growth_to_parameter( + parameter, + nawi=nawi, + start_year=start_year, + end_year=end_year, + ) + return parameters + + class reform(Reform): + def apply(self): + self.modify_parameters(modify_parameters) + + return reform + + +def create_wage_indexed_irs_uprating_reform( + start_year: int = 2035, + end_year: int = 2100, +): + from policyengine_us.model_api import Reform + + def modify_parameters(parameters): + nawi = parameters.gov.ssa.nawi + irs_uprating = parameters.gov.irs.uprating + + for year in range(start_year, end_year + 1): + previous_value = float(irs_uprating(f"{year - 1}-01-01")) + wage_growth = float(nawi(f"{year - 1}-01-01")) / float( + nawi(f"{year - 2}-01-01") + ) + irs_uprating.update( + period=f"year:{year}-01-01:1", + value=previous_value * wage_growth, + ) + return parameters + + class reform(Reform): + def apply(self): + self.modify_parameters(modify_parameters) + + return reform + + def _coerce_h5_path(raw: str) -> Path: path = Path(raw).expanduser() if path.is_dir(): @@ -106,13 +221,28 @@ def _compute_reformed_shares( h5_path: Path, start_year: int, end_year: int, + scenario: str, ) -> dict: from policyengine_us import Microsimulation - reform = create_wage_indexed_brackets_reform( - start_year=start_year, - end_year=end_year, - ) + if scenario == "brackets": + reform = create_wage_indexed_brackets_reform( + start_year=start_year, + end_year=end_year, + ) + elif scenario == "core-thresholds": + reform = create_wage_indexed_core_thresholds_reform( + start_year=start_year, + end_year=end_year, + ) + elif scenario == "irs-uprating": + reform = create_wage_indexed_irs_uprating_reform( + start_year=start_year, + end_year=end_year, + ) + else: + raise ValueError(f"Unknown scenario: {scenario}") + sim = Microsimulation(dataset=str(h5_path), reform=reform) ss_total = float(sim.calculate("social_security").sum()) oasdi_tob = float(sim.calculate("tob_revenue_oasdi").sum()) @@ -125,9 +255,9 @@ def _compute_reformed_shares( def _format_markdown(records: list[dict]) -> str: header = ( - "| Year | Trustees OASDI target | Baseline OASDI | Wage-indexed OASDI | " - "OASDI delta | Baseline combined | Wage-indexed combined |\n" - "| --- | ---: | ---: | ---: | ---: | ---: | ---: |" + "| Year | Scenario | Trustees OASDI target | Baseline OASDI | " + "Reformed OASDI | OASDI delta | Baseline combined | Reformed combined |\n" + "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |" ) rows = [] for record in records: @@ -136,9 +266,10 @@ def _format_markdown(records: list[dict]) -> str: baseline_combined = record["baseline_combined_share_pct"] reformed_combined = record["reformed_combined_share_pct"] rows.append( - "| {year} | {target:.2f}% | {base_o:.2f}% | {reform_o:.2f}% | {delta:+.2f} pp | " + "| {year} | {scenario} | {target:.2f}% | {base_o:.2f}% | {reform_o:.2f}% | {delta:+.2f} pp | " "{base_c:.2f}% | {reform_c:.2f}% |".format( year=record["year"], + scenario=record["scenario"], target=record["target_oasdi_share_pct"], base_o=baseline_oasdi, reform_o=reformed_oasdi, @@ -178,8 +309,10 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--end-year", type=int, - default=2100, - help="Last year to extend the wage-indexed bracket sensitivity through.", + help=( + "Last year to extend the wage-indexed sensitivity through. " + "Defaults to the maximum year among the input H5s." + ), ) parser.add_argument( "--format", @@ -187,6 +320,16 @@ def parse_args() -> argparse.Namespace: default="markdown", help="Output format.", ) + parser.add_argument( + "--scenario", + choices=("brackets", "core-thresholds", "irs-uprating"), + default="brackets", + help=( + "Tax-side sensitivity to run: wage-index only ordinary bracket " + "thresholds, wage-index a core threshold set, or wage-index the " + "full IRS uprating path." + ), + ) return parser.parse_args() @@ -196,22 +339,25 @@ def main() -> int: sys.path.insert(0, str(Path(args.policyengine_us_path).expanduser())) h5_paths = [_coerce_h5_path(path) for path in args.paths] + end_year = args.end_year or max(int(path.stem) for path in h5_paths) records = [] for h5_path in h5_paths: metadata = _load_metadata(h5_path) baseline = _baseline_record(h5_path, metadata) print( - f"[{baseline['year']}] benchmarking wage-indexed federal brackets on {h5_path}", + f"[{baseline['year']}] benchmarking {args.scenario} on {h5_path}", file=sys.stderr, flush=True, ) reformed = _compute_reformed_shares( h5_path, start_year=args.start_year, - end_year=args.end_year, + end_year=end_year, + scenario=args.scenario, ) baseline.update(reformed) + baseline["scenario"] = args.scenario records.append(baseline) records.sort(key=lambda record: record["year"]) From 7ffc8ed57154b2ead29f124bbe9f985de0d4fc9d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 2 Apr 2026 08:29:36 -0400 Subject: [PATCH 28/31] Benchmark full IRS uprating for TOB --- .../cps/long_term/TOB_ALIGNMENT_NOTE.md | 26 ++++++++++++ .../benchmark_trustees_bracket_indexing.py | 42 ++++++++++++------- 2 files changed, 54 insertions(+), 14 deletions(-) diff --git a/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md b/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md index 06c01f33a..092925560 100644 --- a/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md +++ b/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md @@ -121,6 +121,32 @@ Interpretation: beneficiary income mix, filing composition, or other Treasury-ratio modeling differences. +## Full IRS-uprating upper bound + +Finally, we ran an upper-bound sensitivity that rewrites every materialized IRS +parameter leaf that currently inherits from `gov.irs.uprating`, replacing +post-`2034` `C-CPI-U` growth with `NAWI` growth. + +This is broader than the public Trustees text justifies, but it provides a +useful ceiling on how much of the TOB gap could plausibly be explained by the +IRS uprating family alone. + +| Year | Baseline OASDI | Full IRS-uprating OASDI | Trustees target | Remaining gap | +| --- | ---: | ---: | ---: | ---: | +| 2075 | 9.43% | 7.46% | 6.01% | +1.45 pp | +| 2090 | 10.52% | 7.17% | 6.08% | +1.09 pp | +| 2100 | 11.16% | 8.16% | 6.10% | +2.06 pp | + +Interpretation: + +- The full IRS-uprating upper bound is only slightly lower than the narrower + core-threshold bundle. +- That implies most of the tax-side movement is already coming from the core + federal threshold families, not from the rest of the CPI-uprated IRS + parameter tree. +- Even under this broad upper bound, the model still remains above the + Trustees OASDI-only TOB path, especially in `2100`. + ## DYNASIM public benchmark Urban's 2024 appendix says DYNASIM's revenue from taxing Social Security diff --git a/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py b/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py index 95a11501c..23da92226 100644 --- a/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py +++ b/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py @@ -35,7 +35,15 @@ def round_amount(amount: float, rounding: dict | None) -> float: raise ValueError(f"Unsupported rounding type: {rounding_type}") -def _iter_updatable_parameters(root) -> list: +def _uprating_parameter_name(parameter) -> str | None: + metadata = getattr(parameter, "metadata", {}) + uprating = metadata.get("uprating") + if isinstance(uprating, dict): + return uprating.get("parameter") + return uprating + + +def _iter_updatable_parameters(root, uprating_parameter: str | None = None) -> list: candidates = [root] if hasattr(root, "get_descendants"): candidates.extend(root.get_descendants()) @@ -44,9 +52,12 @@ def _iter_updatable_parameters(root) -> list: for candidate in candidates: if candidate.__class__.__name__ != "Parameter": continue - uprating = getattr(candidate, "metadata", {}).get("uprating") - if uprating is not None: - result.append(candidate) + uprating_name = _uprating_parameter_name(candidate) + if uprating_name is None: + continue + if uprating_parameter is not None and uprating_name != uprating_parameter: + continue + result.append(candidate) return result @@ -152,16 +163,19 @@ def create_wage_indexed_irs_uprating_reform( def modify_parameters(parameters): nawi = parameters.gov.ssa.nawi - irs_uprating = parameters.gov.irs.uprating - - for year in range(start_year, end_year + 1): - previous_value = float(irs_uprating(f"{year - 1}-01-01")) - wage_growth = float(nawi(f"{year - 1}-01-01")) / float( - nawi(f"{year - 2}-01-01") - ) - irs_uprating.update( - period=f"year:{year}-01-01:1", - value=previous_value * wage_growth, + seen = set() + for parameter in _iter_updatable_parameters( + parameters.gov.irs, + uprating_parameter="gov.irs.uprating", + ): + if parameter.name in seen: + continue + seen.add(parameter.name) + _apply_wage_growth_to_parameter( + parameter, + nawi=nawi, + start_year=start_year, + end_year=end_year, ) return parameters From 9f46122e57945405dc5be4d0ca11e91e8b2a92df Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 2 Apr 2026 08:47:22 -0400 Subject: [PATCH 29/31] Adopt core-threshold TOB baseline --- .../cps/long_term/ASSUMPTION_COMPARISON.md | 35 +- .../datasets/cps/long_term/README.md | 52 +- .../cps/long_term/TOB_ALIGNMENT_NOTE.md | 33 +- .../benchmark_trustees_bracket_indexing.py | 151 +----- .../cps/long_term/calibration_artifacts.py | 21 +- .../cps/long_term/calibration_profiles.py | 12 +- .../cps/long_term/compare_tob_shares.py | 10 +- .../cps/long_term/projection_utils.py | 42 +- .../prototype_synthetic_2100_support.py | 383 ++++++++++----- .../cps/long_term/run_household_projection.py | 445 +++++++++++++++--- .../datasets/cps/long_term/tax_assumptions.py | 178 +++++++ .../test_long_term_calibration_contract.py | 181 ++++++- 12 files changed, 1185 insertions(+), 358 deletions(-) create mode 100644 policyengine_us_data/datasets/cps/long_term/tax_assumptions.py diff --git a/policyengine_us_data/datasets/cps/long_term/ASSUMPTION_COMPARISON.md b/policyengine_us_data/datasets/cps/long_term/ASSUMPTION_COMPARISON.md index 4a9764488..8a1278f8e 100644 --- a/policyengine_us_data/datasets/cps/long_term/ASSUMPTION_COMPARISON.md +++ b/policyengine_us_data/datasets/cps/long_term/ASSUMPTION_COMPARISON.md @@ -3,24 +3,41 @@ This note distinguishes between: - hard microsimulation calibration targets, which directly shape household weights -- post-calibration benchmarks, which are compared against published source packages without being forced into the weights +- tax-side assumptions used to make those targets more comparable to the public Trustees/OACT methodology -The split is intentional. Long-run Taxation of Benefits (TOB) depends on tax-side assumptions that are not fully published as a household-level rule schedule in the Trustees/OACT materials. +The current long-run baseline now adopts a named tax-side assumption, +`trustees-core-thresholds-v1`, before hard-targeting TOB. | Component | Current `policyengine-us-data` approach | Trustees / OACT published approach | Calibration use | | --- | --- | --- | --- | | Population by age | SSA single-year age projections | SSA single-year age projections | Hard target | | OASDI benefits | Named long-term target source package | Trustees or OACT-patched annual OASDI path | Hard target | | Taxable payroll | Named long-term target source package | Trustees annual taxable payroll path | Hard target | -| Social Security benefit-tax thresholds | Literal current-law statutory thresholds remain fixed in nominal dollars | Trustees also describe statutory thresholds as remaining fixed in nominal dollars | Not separately targeted | -| Federal income-tax brackets | PolicyEngine tax simulation using its own income-tax parameter uprating path | Trustees assume periodic future bracket adjustments; after the tenth year, tax brackets are assumed to rise with average wages to avoid indefinite bracket creep | Not hard-targeted | -| OASDI TOB | Computed from the tax microsimulation and compared to the selected target source | Trustees/OACT publish annual revenue paths or ratios, but not a full public household-level micro rule schedule | Post-calibration benchmark | -| HI TOB | Computed from the tax microsimulation and compared to the selected target source | Trustees publish current-law HI TOB path; OACT OBBBA updates do not currently provide a full public annual HI replacement series | Post-calibration benchmark | +| Social Security benefit-tax thresholds | Literal current-law statutory thresholds remain fixed in nominal dollars | Trustees also describe the statutory `$25k/$32k/$0` and `$34k/$44k` thresholds as remaining fixed in nominal dollars | Not separately targeted | +| Federal income-tax brackets | Core ordinary thresholds are wage-indexed after `2034` via `trustees-core-thresholds-v1` | Trustees assume periodic future bracket adjustments; after the tenth projection year, ordinary federal income-tax brackets are assumed to rise with average wages to avoid indefinite bracket creep | Tax-side assumption | +| Standard deduction / aged-blind addition / capital gains thresholds / AMT thresholds | Included in the same `trustees-core-thresholds-v1` bundle | Not parameterized publicly line-by-line, but these are the main additional federal thresholds most likely to affect long-run TOB | Tax-side assumption | +| OASDI TOB | Computed under the core-threshold tax assumption and targeted in `ss-payroll-tob` profiles | Trustees/OACT publish annual revenue paths or ratios, but not a full public household-level micro rule schedule | Hard target | +| HI TOB | Computed under the core-threshold tax assumption and targeted in `ss-payroll-tob` profiles | Trustees publish current-law HI TOB path; OACT OBBBA updates do not currently provide a full public annual HI replacement series | Hard target | | OBBBA OASDI update | Available through named target source `oact_2025_08_05_provisional` | August 5, 2025 OACT letter provides annual OASDI changes through 2100 | Benchmark / target-source input | | OBBBA HI update | Provisional bridge only in named target source | No equivalent full public annual HI replacement path located yet | Benchmark only | ## Practical interpretation -- `ss-payroll` is the core hard-target profile. -- `ss-payroll-tob` now means: calibrate on age + OASDI benefits + taxable payroll, then benchmark TOB against the selected source package. -- TOB remains important for model comparison, but it is no longer treated as a weight-identifying target in the long-run microsimulation contract. +- `ss-payroll` remains the core non-TOB hard-target profile. +- `ss-payroll-tob` now means: calibrate on age + OASDI benefits + taxable payroll + TOB under `trustees-core-thresholds-v1`. +- The core-threshold bundle is a best-public approximation, not a literal public Trustees rules schedule. +- Trustees-consistent long-run TOB requires keeping two different tax-side ideas separate: + - the Social Security benefit-tax thresholds remain fixed in nominal dollars + - ordinary federal income-tax brackets are assumed to rise with average wages after the tenth projection year + +## Primary-source references + +- [SSA 2025 Trustees Report, V.C.7](https://www.ssa.gov/oact/tr/2025/V_C_prog.html) + - States that the law specifies fixed threshold amounts for taxation of Social Security benefits and that those thresholds remain constant in future years. + - Also states that, after the tenth year of the projection period, income-tax brackets are assumed to rise with average wages rather than with `C-CPI-U`. +- [26 U.S.C. § 86](https://www.law.cornell.edu/uscode/text/26/86) + - Statutory basis for the Social Security benefit-tax threshold structure. +- [SSA 2025 Trustees Report, Table VI.G6](https://www.ssa.gov/OACT/TR/2025/VI_G3_OASDHI_dollars.html) + - Published annual average wage index path through `2100`. +- [42 U.S.C. § 430](https://www.law.cornell.edu/uscode/text/42/430) and [20 CFR § 404.1048](https://www.law.cornell.edu/cfr/text/20/404.1048) + - Statutory and regulatory basis for deriving the Social Security contribution and benefit base from the wage index. diff --git a/policyengine_us_data/datasets/cps/long_term/README.md b/policyengine_us_data/datasets/cps/long_term/README.md index 60c21e660..e88d4d6fb 100644 --- a/policyengine_us_data/datasets/cps/long_term/README.md +++ b/policyengine_us_data/datasets/cps/long_term/README.md @@ -6,7 +6,7 @@ Run projections using `run_household_projection.py`: ```bash -# Recommended: named profile with TOB benchmarked post-calibration +# Recommended: named profile with core-threshold tax assumption and TOB targeted python run_household_projection.py 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --save-h5 # Experimental: donor-backed late-year support augmentation for tail-year runs @@ -15,6 +15,9 @@ python run_household_projection.py 2075 2100 --profile ss-payroll-tob --target-s # Experimental: role-based donor composites assembled into late-year support python run_household_projection.py 2075 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --support-augmentation-profile donor-backed-composite-v1 --support-augmentation-target-year 2100 --allow-validation-failures +# Experimental: target-year blueprint calibration over donor-composite support +python run_household_projection.py 2100 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --support-augmentation-profile donor-backed-composite-v1 --support-augmentation-target-year 2100 --support-augmentation-blueprint-base-weight-scale 0.5 --allow-validation-failures --save-h5 + # IPF with only age distribution constraints (faster, less accurate) python run_household_projection.py 2050 --profile age-only @@ -26,14 +29,17 @@ python run_household_projection.py 2100 --profile ss - `END_YEAR`: Target year for projection (default: 2035) - `--profile`: Named calibration contract. Recommended over legacy flags. - `--target-source`: Named long-term target source package. +- `--tax-assumption`: Long-run federal tax assumption. Defaults to `trustees-core-thresholds-v1`; use `current-law-literal` to opt out. - `--output-dir`: Output directory for generated H5 files and metadata sidecars. - `--support-augmentation-profile`: Experimental late-year support expansion mode. Currently supports `donor-backed-synthetic-v1` and `donor-backed-composite-v1`. - `--support-augmentation-target-year`: Extreme year used to build the donor-backed supplement (defaults to `END_YEAR`). +- `--support-augmentation-align-to-run-year`: Rebuild the donor-backed supplement separately for each run year instead of reusing one target-year support snapshot. - `--support-augmentation-start-year`: Earliest run year allowed for augmentation (defaults to `2075`). - `--support-augmentation-top-n-targets`: Number of dominant synthetic target types to map back to real donors (default `20`). - `--support-augmentation-donors-per-target`: Number of nearest real donor tax units per synthetic target (default `5`). - `--support-augmentation-max-distance`: Maximum donor-match distance retained for cloning (default `3.0`). - `--support-augmentation-clone-weight-scale`: Baseline weight multiplier applied to each donor-backed clone (default `0.1`). +- `--support-augmentation-blueprint-base-weight-scale`: When donor-composite augmentation is active at its target year, scales the original household priors before replacing clone priors with synthetic blueprint shares (default `0.5`). - `--greg`: Use GREG calibration instead of IPF - `--use-ss`: Include Social Security benefit totals as calibration target (requires `--greg`) - `--use-payroll`: Include taxable payroll totals as calibration target (requires `--greg`) @@ -44,8 +50,8 @@ python run_household_projection.py 2100 --profile ss - `age-only`: IPF age-only calibration - `ss`: positive entropy calibration with age + Social Security - `ss-payroll`: positive entropy calibration with age + Social Security + taxable payroll -- `ss-payroll-tob`: positive entropy calibration with age + Social Security + taxable payroll, with TOB benchmarked after calibration -- `ss-payroll-tob-h6`: positive entropy calibration with age + Social Security + taxable payroll + H6, with TOB benchmarked after calibration +- `ss-payroll-tob`: positive entropy calibration with age + Social Security + taxable payroll + TOB under the long-run core-threshold tax assumption +- `ss-payroll-tob-h6`: positive entropy calibration with age + Social Security + taxable payroll + TOB + H6 under the long-run core-threshold tax assumption **Validation contract:** - Economic-targeted profiles no longer silently pretend an IPF fallback is equivalent to GREG. @@ -56,7 +62,15 @@ python run_household_projection.py 2100 --profile ss profile/base dataset contract for the full artifact set. - Profiles validate achieved constraint errors before writing output. - Experimental donor-backed augmentation is stamped into each year sidecar and the directory manifest via `support_augmentation`. +- The active long-run tax assumption is stamped into each year sidecar and the directory manifest via `tax_assumption`. - Donor-backed runs now also write a shared `support_augmentation_report.json` artifact with per-clone provenance so late-year translation failures can be inspected directly. +- Long-run payroll calibration now guards against a flat Social Security wage base after 2035. If `policyengine-us` is missing the NAWI / payroll-cap extension, late-year payroll runs fail fast instead of silently mis-targeting taxable payroll. +- Trustees/OACT tax-side assumptions are documented in [ASSUMPTION_COMPARISON.md](./ASSUMPTION_COMPARISON.md). The active long-run baseline adopts a core-threshold bundle: + - Social Security benefit-tax thresholds remain fixed in nominal dollars under Trustees current law. + - Core ordinary federal thresholds are assumed to rise with average wages after the tenth projection year. +- The corrected apples-to-apples TOB share comparison is documented in + [TOB_ALIGNMENT_NOTE.md](./TOB_ALIGNMENT_NOTE.md), with a small reproduction + script in [compare_tob_shares.py](./compare_tob_shares.py). **Estimated runtime:** ~2 minutes/year without `--save-h5`, ~3 minutes/year with `--save-h5` @@ -75,20 +89,35 @@ python run_household_projection.py 2100 --profile ss - Uses dual optimization to minimize divergence from baseline weights - **Recommended** for publishable long-term projections +**Long-Run Core-Threshold Tax Assumption** +- Default long-run tax assumption in this runner +- Keeps Social Security benefit-tax thresholds fixed +- Wage-indexes a core set of federal thresholds after `2034`: + - ordinary income-tax brackets + - standard deduction + - aged/blind additional standard deduction + - capital-gains thresholds + - AMT thresholds / exemptions +- Intended as the best public Trustees approximation before TOB is hard-targeted again + **Donor-Backed Late-Year Support Augmentation** - Experimental late-tail option for `2075+` runs - Uses the `2100` synthetic-support prototype to identify dominant missing household types - Maps those synthetic targets back to nearest real 2024 donor tax units - Clones and perturbs the donor tax units to create a small augmented support without replacing the base CPS sample - Intended to test whether donor-backed synthetic support improves late-year microsim feasibility without resorting to fully free synthetic records -- Current status: integrated into the runner and fully auditable in metadata, but still diagnostic. The first `2100` end-to-end run did not materially improve the late-tail calibration frontier or support-concentration metrics. After adding clone-provenance diagnostics and donor-specific inverse uprating, the realized clone households now match their intended `2100` age/SS/payroll targets closely; the remaining blocker is how much synthetic support is injected, not whether the actual-row translation works. +- Current status: still diagnostic. The simple nearest-neighbor donor supplement does not materially improve the late-tail fit once the calibration uses SSA taxable payroll rather than uncapped wages. **Role-Based Donor Composites** - Experimental structural extension of the donor-backed approach - Recombines older-beneficiary donors, payroll-rich worker donors, and dependent structure into synthetic household candidates before assembling actual augmented rows -- In the synthetic support lab, this materially improves the `2100` exact-fit basis: exact entropy fit with `360` positive candidates, ESS about `95.8`, and top-10 weight share about `25.6%` - The actual-row augmented dataset builder is now available in the runner as `donor-backed-composite-v1` -- Current runner status: the first full `2100` end-to-end run with the default composite supplement added `270` structural clones and modestly improved support concentration (`ESS 11.4 -> 13.3`, top-10 share `84.9% -> 79.4%`, lower TOB overshoot), but it did not yet move the main `SS + payroll` frontier off the current `-33.8% / -35.0%` late-tail bound +- Current status: + - Fixing the long-run payroll-cap bug in `policyengine-us` changed the picture materially. With the correct SSA wage base extended through `2100`, the donor-composite synthetic support is exact-feasible and dense at the archetype level. + - The runner now supports a target-year calibration blueprint for donor-composite augmentation. At the augmentation target year, it can calibrate against the exact clone blueprints and synthetic prior shares while still auditing the realized rows. + - In the current `2100` probe, that blueprint path gets actual age + SS + taxable payroll very close while keeping support quality in range: with `--support-augmentation-blueprint-base-weight-scale 0.5`, actual payroll miss is about `-0.86%`, ESS about `102.5`, top-10 weight share about `24.4%`, and top-100 share about `68.4%`. + - The runner now also has a dynamic mode, `--support-augmentation-align-to-run-year`, that rebuilds donor-composite support for each run year and writes per-year augmentation reports. + - This is still experimental. The blueprint path is now structurally capable of handling year-specific support, but the full `2075-2100` production sweep still needs runtime tuning and caching work. **GREG (Generalized Regression Estimator)** - Legacy linear calibration path retained for explicit flag-based runs @@ -111,13 +140,15 @@ python run_household_projection.py 2100 --profile ss - W-2 wages capped at wage base + SE income within remaining cap room - Calculated as: `taxable_earnings_for_social_security` + `social_security_taxable_self_employment_income` - Source: selected long-term target source package + - Guardrail: the runner checks that the Social Security taxable earnings cap continues to rise after 2035. A flat cap indicates an invalid `policyengine-us` parameter baseline for long-run payroll work. -4. **TOB Revenue** (`--use-tob`, legacy hard-target mode only) +4. **TOB Revenue** (`ss-payroll-tob`, `ss-payroll-tob-h6`, or legacy `--use-tob`) - Taxation of Benefits revenue for OASDI and Medicare HI trust funds - OASDI: `tob_revenue_oasdi` (tier 1 taxation, 0-50% of benefits) - HI: `tob_revenue_medicare_hi` (tier 2 taxation, 50-85% of benefits) - Source: selected long-term target source package - - Recommended usage: benchmark after calibration rather than use as a hard weight target + - In the current branch contract, TOB is hard-targeted under the long-run core-threshold tax assumption rather than under literal CPI-style current-law bracket indexing. + - Primary-source note: Trustees do not appear to model long-run TOB by indexing the SS-specific `$25k/$32k/$0` and `$34k/$44k` thresholds. Those remain fixed under current law; the long-run divergence comes from broader tax-side assumptions, including ordinary bracket treatment after the tenth projection year. --- @@ -127,6 +158,7 @@ python run_household_projection.py 2100 --profile ss - URL: https://www.ssa.gov/OACT/TR/2025/ - File: `SingleYearTRTables_TR2025.xlsx` - Tables: IV.B2 (OASDI TOB % of taxable payroll), VI.G6 (taxable payroll in billions), VI.G9 (OASDI costs) +- Program assumptions: [V.C.7](https://www.ssa.gov/oact/tr/2025/V_C_prog.html) documents fixed Social Security benefit-tax thresholds and the long-run ordinary-bracket assumption. **CMS 2025 Medicare Trustees Report** - URL: https://www.cms.gov/data-research/statistics-trends-and-reports/trustees-report-trust-funds @@ -173,11 +205,13 @@ For each projection year (2025-2100): 2. **Uprate variables** - PolicyEngine automatically uprates income, thresholds, etc. to target year 3. **Calculate values** - Income tax, Social Security, taxable payroll at household level 4. **Calibrate weights** - Adjust household weights to match SSA demographic/economic targets -5. **Benchmark TOB** - Compare modeled OASDI/HI TOB to the selected target source without forcing it into the weights +5. **Target or benchmark TOB** - Under `ss-payroll-tob`, match modeled OASDI/HI TOB to the selected target source using the core-threshold tax assumption 6. **Aggregate results** - Apply calibrated weights to calculate national totals When donor-backed augmentation is enabled, step 1 uses the original 2024 CPS support and step 2 inserts a tagged late-year supplement derived from nearest real donors before the calibration loop begins. The underlying base dataset path remains unchanged in metadata; the augmentation details are recorded separately in `support_augmentation`, and the full augmentation build report is written once per output directory to `support_augmentation_report.json`. +When donor-composite augmentation is enabled and the run year equals the augmentation target year, the runner can also replace the clone rows' calibration constraints with their exact synthetic blueprint values and use the synthetic-support solution as clone priors. The calibration audit still reports achieved constraints on the realized rows, so any blueprint-to-row translation gap remains visible in metadata. + To compare the intended clone targets with the realized output H5, run: ```bash diff --git a/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md b/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md index 092925560..ede8c578a 100644 --- a/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md +++ b/policyengine_us_data/datasets/cps/long_term/TOB_ALIGNMENT_NOTE.md @@ -29,6 +29,33 @@ not the latter. `trustees_2025_current_law.csv` target package imply OASDI-only TOB shares of about `6.0%` to `6.1%` of OASDI benefits in the late horizon. +## Current contract decision + +As of `2026-04-02`, the branch adopts the `trustees-core-thresholds-v1` +tax-side assumption for long-run TOB work and re-enables TOB as a hard target +in the `ss-payroll-tob` profiles. + +That assumption: + +- keeps the Social Security benefit-tax thresholds fixed +- wage-indexes a core federal threshold bundle after `2034` + - ordinary income-tax brackets + - standard deduction + - aged/blind additional standard deduction + - capital-gains thresholds + - AMT thresholds / exemptions + +On a `2100` smoke run with donor-composite support augmentation and the local +`policyengine-us` wage-base fix, this produces a validation-passing artifact +with: + +- `ss_total`: essentially exact +- `oasdi_tob`: essentially exact +- `hi_tob`: essentially exact +- `payroll_total`: `-2.95%` +- `ESS`: `173.1` +- `top-10 weight share`: `17.7%` + ## Current branch-local comparison The table below uses one-year probe outputs produced on `2026-04-02` with: @@ -185,11 +212,11 @@ uv run python policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bra --policyengine-us-path /path/to/local/policyengine-us ``` -The script expects metadata sidecars containing: +The scripts expect metadata sidecars containing: - `calibration_audit.constraints.ss_total` -- `calibration_audit.benchmarks.oasdi_tob` -- `calibration_audit.benchmarks.hi_tob` +- `calibration_audit.constraints.oasdi_tob` or `calibration_audit.benchmarks.oasdi_tob` +- `calibration_audit.constraints.hi_tob` or `calibration_audit.benchmarks.hi_tob` To regenerate the underlying sidecars from scratch, first ensure that `policyengine-us` includes the Social Security wage-base extension from diff --git a/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py b/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py index 23da92226..5d44edbe6 100644 --- a/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py +++ b/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py @@ -6,6 +6,17 @@ import sys from pathlib import Path +try: + from .tax_assumptions import ( + create_wage_indexed_core_thresholds_reform, + create_wage_indexed_full_irs_uprating_reform, + ) +except ImportError: # pragma: no cover - script execution fallback + from tax_assumptions import ( + create_wage_indexed_core_thresholds_reform, + create_wage_indexed_full_irs_uprating_reform, + ) + FILING_STATUSES = [ "SINGLE", @@ -20,64 +31,6 @@ def round_down(amount: float, interval: float) -> float: return math.floor(amount / interval) * interval -def round_amount(amount: float, rounding: dict | None) -> float: - if not rounding: - return amount - - interval = float(rounding["interval"]) - rounding_type = rounding["type"] - - if rounding_type == "downwards": - return math.floor(amount / interval) * interval - if rounding_type == "nearest": - return math.floor(amount / interval + 0.5) * interval - - raise ValueError(f"Unsupported rounding type: {rounding_type}") - - -def _uprating_parameter_name(parameter) -> str | None: - metadata = getattr(parameter, "metadata", {}) - uprating = metadata.get("uprating") - if isinstance(uprating, dict): - return uprating.get("parameter") - return uprating - - -def _iter_updatable_parameters(root, uprating_parameter: str | None = None) -> list: - candidates = [root] - if hasattr(root, "get_descendants"): - candidates.extend(root.get_descendants()) - - result = [] - for candidate in candidates: - if candidate.__class__.__name__ != "Parameter": - continue - uprating_name = _uprating_parameter_name(candidate) - if uprating_name is None: - continue - if uprating_parameter is not None and uprating_name != uprating_parameter: - continue - result.append(candidate) - return result - - -def _apply_wage_growth_to_parameter(parameter, nawi, start_year: int, end_year: int): - metadata = getattr(parameter, "metadata", {}) - uprating = metadata.get("uprating") - rounding = uprating.get("rounding") if isinstance(uprating, dict) else None - - for year in range(start_year, end_year + 1): - previous_value = float(parameter(f"{year - 1}-01-01")) - wage_growth = float(nawi(f"{year - 1}-01-01")) / float( - nawi(f"{year - 2}-01-01") - ) - updated_value = round_amount(previous_value * wage_growth, rounding) - parameter.update( - period=f"year:{year}-01-01:1", - value=updated_value, - ) - - def create_wage_indexed_brackets_reform( start_year: int = 2035, end_year: int = 2100, @@ -113,79 +66,6 @@ def apply(self): self.modify_parameters(modify_parameters) return reform - - -def create_wage_indexed_core_thresholds_reform( - start_year: int = 2035, - end_year: int = 2100, -): - from policyengine_us.model_api import Reform - - def modify_parameters(parameters): - nawi = parameters.gov.ssa.nawi - roots = [ - parameters.gov.irs.income.bracket.thresholds, - parameters.gov.irs.deductions.standard.amount, - parameters.gov.irs.deductions.standard.aged_or_blind.amount, - parameters.gov.irs.capital_gains.thresholds, - parameters.gov.irs.income.amt.brackets, - parameters.gov.irs.income.amt.exemption.amount, - parameters.gov.irs.income.amt.exemption.phase_out.start, - parameters.gov.irs.income.amt.exemption.separate_limit, - ] - - seen = set() - for root in roots: - for parameter in _iter_updatable_parameters(root): - if parameter.name in seen: - continue - seen.add(parameter.name) - _apply_wage_growth_to_parameter( - parameter, - nawi=nawi, - start_year=start_year, - end_year=end_year, - ) - return parameters - - class reform(Reform): - def apply(self): - self.modify_parameters(modify_parameters) - - return reform - - -def create_wage_indexed_irs_uprating_reform( - start_year: int = 2035, - end_year: int = 2100, -): - from policyengine_us.model_api import Reform - - def modify_parameters(parameters): - nawi = parameters.gov.ssa.nawi - seen = set() - for parameter in _iter_updatable_parameters( - parameters.gov.irs, - uprating_parameter="gov.irs.uprating", - ): - if parameter.name in seen: - continue - seen.add(parameter.name) - _apply_wage_growth_to_parameter( - parameter, - nawi=nawi, - start_year=start_year, - end_year=end_year, - ) - return parameters - - class reform(Reform): - def apply(self): - self.modify_parameters(modify_parameters) - - return reform - - def _coerce_h5_path(raw: str) -> Path: path = Path(raw).expanduser() if path.is_dir(): @@ -219,9 +99,10 @@ def _baseline_record(h5_path: Path, metadata: dict | None) -> dict: audit = metadata["calibration_audit"] ss_total = float(audit["constraints"]["ss_total"]["achieved"]) ss_target = float(audit["constraints"]["ss_total"]["target"]) - oasdi_actual = float(audit["benchmarks"]["oasdi_tob"]["achieved"]) - hi_actual = float(audit["benchmarks"]["hi_tob"]["achieved"]) - oasdi_target = float(audit["benchmarks"]["oasdi_tob"]["target"]) + tob_section = audit.get("benchmarks") or audit.get("constraints") + oasdi_actual = float(tob_section["oasdi_tob"]["achieved"]) + hi_actual = float(tob_section["hi_tob"]["achieved"]) + oasdi_target = float(tob_section["oasdi_tob"]["target"]) return { "year": int(metadata["year"]), @@ -250,7 +131,7 @@ def _compute_reformed_shares( end_year=end_year, ) elif scenario == "irs-uprating": - reform = create_wage_indexed_irs_uprating_reform( + reform = create_wage_indexed_full_irs_uprating_reform( start_year=start_year, end_year=end_year, ) diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py index ece85be86..610a7a83e 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py @@ -103,6 +103,7 @@ def write_year_metadata( profile: dict[str, Any], calibration_audit: dict[str, Any], target_source: dict[str, Any] | None = None, + tax_assumption: dict[str, Any] | None = None, support_augmentation: dict[str, Any] | None = None, ) -> Path: metadata = { @@ -114,6 +115,8 @@ def write_year_metadata( } if target_source is not None: metadata["target_source"] = target_source + if tax_assumption is not None: + metadata["tax_assumption"] = tax_assumption if support_augmentation is not None: metadata["support_augmentation"] = support_augmentation metadata = normalize_metadata(metadata) @@ -126,11 +129,14 @@ def write_year_metadata( def write_support_augmentation_report( - output_dir: str | Path, report: dict[str, Any] + output_dir: str | Path, + report: dict[str, Any], + *, + filename: str = SUPPORT_AUGMENTATION_REPORT_FILENAME, ) -> Path: output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) - report_path = output_dir / SUPPORT_AUGMENTATION_REPORT_FILENAME + report_path = output_dir / filename report_path.write_text( json.dumps(json.loads(json.dumps(report)), indent=2, sort_keys=True) + "\n", encoding="utf-8", @@ -148,12 +154,14 @@ def update_dataset_manifest( profile: dict[str, Any], calibration_audit: dict[str, Any], target_source: dict[str, Any] | None = None, + tax_assumption: dict[str, Any] | None = None, support_augmentation: dict[str, Any] | None = None, ) -> Path: output_dir = Path(output_dir) manifest_path = output_dir / MANIFEST_FILENAME profile = json.loads(json.dumps(profile)) target_source = json.loads(json.dumps(target_source)) + tax_assumption = json.loads(json.dumps(tax_assumption)) support_augmentation = json.loads(json.dumps(support_augmentation)) if manifest_path.exists(): @@ -165,6 +173,7 @@ def update_dataset_manifest( "base_dataset_path": base_dataset_path, "profile": profile, "target_source": target_source, + "tax_assumption": tax_assumption, "support_augmentation": support_augmentation, "years": [], "datasets": {}, @@ -195,6 +204,13 @@ def update_dataset_manifest( "Output directory already contains a different target source: " f"{manifest.get('target_source')} != {target_source}" ) + if manifest.get("tax_assumption") is None and tax_assumption is not None: + manifest["tax_assumption"] = tax_assumption + elif manifest.get("tax_assumption") != tax_assumption: + raise ValueError( + "Output directory already contains a different tax assumption: " + f"{manifest.get('tax_assumption')} != {tax_assumption}" + ) if ( manifest.get("support_augmentation") is None and support_augmentation is not None @@ -281,6 +297,7 @@ def rebuild_dataset_manifest_with_target_source( profile=metadata["profile"], calibration_audit=metadata["calibration_audit"], target_source=metadata.get("target_source"), + tax_assumption=metadata.get("tax_assumption"), support_augmentation=metadata.get("support_augmentation"), ) diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py index a0a49298b..b093bc882 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py @@ -157,14 +157,14 @@ def to_dict(self) -> dict[str, Any]: ), "ss-payroll-tob": CalibrationProfile( name="ss-payroll-tob", - description="Age, Social Security, and taxable payroll using positive entropy calibration, with TOB benchmarked post-calibration against the selected long-term target source.", + description="Age, Social Security, taxable payroll, and TOB using positive entropy calibration under the long-run core-threshold tax assumption.", calibration_method="entropy", use_greg=False, use_ss=True, use_payroll=True, use_h6_reform=False, - use_tob=False, - benchmark_tob=True, + use_tob=True, + benchmark_tob=False, max_negative_weight_pct=0.0, min_positive_household_count=1000, min_effective_sample_size=75.0, @@ -174,14 +174,14 @@ def to_dict(self) -> dict[str, Any]: ), "ss-payroll-tob-h6": CalibrationProfile( name="ss-payroll-tob-h6", - description="Age, Social Security, taxable payroll, and H6 using positive entropy calibration, with TOB benchmarked post-calibration against the selected long-term target source.", + description="Age, Social Security, taxable payroll, TOB, and H6 using positive entropy calibration under the long-run core-threshold tax assumption.", calibration_method="entropy", use_greg=False, use_ss=True, use_payroll=True, use_h6_reform=True, - use_tob=False, - benchmark_tob=True, + use_tob=True, + benchmark_tob=False, max_negative_weight_pct=0.0, min_positive_household_count=1000, min_effective_sample_size=75.0, diff --git a/policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py b/policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py index e3a7e6a4f..9ab4104e8 100644 --- a/policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py +++ b/policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py @@ -26,14 +26,14 @@ def _load_record(path: Path) -> dict: metadata = json.loads(path.read_text(encoding="utf-8")) audit = metadata["calibration_audit"] constraints = audit["constraints"] - benchmarks = audit["benchmarks"] + tob_section = audit.get("benchmarks") or audit.get("constraints") ss_actual = float(constraints["ss_total"]["achieved"]) ss_target = float(constraints["ss_total"]["target"]) - oasdi_actual = float(benchmarks["oasdi_tob"]["achieved"]) - oasdi_target = float(benchmarks["oasdi_tob"]["target"]) - hi_actual = float(benchmarks["hi_tob"]["achieved"]) - hi_target = float(benchmarks["hi_tob"]["target"]) + oasdi_actual = float(tob_section["oasdi_tob"]["achieved"]) + oasdi_target = float(tob_section["oasdi_tob"]["target"]) + hi_actual = float(tob_section["hi_tob"]["achieved"]) + hi_target = float(tob_section["hi_tob"]["target"]) combined_actual = oasdi_actual + hi_actual combined_target = oasdi_target + hi_target diff --git a/policyengine_us_data/datasets/cps/long_term/projection_utils.py b/policyengine_us_data/datasets/cps/long_term/projection_utils.py index 7c75122f6..7c28005f6 100644 --- a/policyengine_us_data/datasets/cps/long_term/projection_utils.py +++ b/policyengine_us_data/datasets/cps/long_term/projection_utils.py @@ -8,6 +8,36 @@ from policyengine_core.data.dataset import Dataset +def validate_projected_social_security_cap( + parameter_accessor, + year: int, + *, + reference_year: int = 2035, +) -> float: + """ + Ensure the Social Security taxable earnings cap keeps growing beyond the + last explicitly projected year. + + The long-run calibration and diagnostics use taxable payroll targets + through 2100. If the payroll cap flattens after the reference year, the + late-year taxable payroll problem becomes mechanically distorted. + """ + current_cap = float( + parameter_accessor(year).gov.irs.payroll.social_security.cap + ) + reference_cap = float( + parameter_accessor(reference_year).gov.irs.payroll.social_security.cap + ) + if year > reference_year and current_cap <= reference_cap * (1 + 1e-12): + raise RuntimeError( + "Social Security payroll cap is flat after " + f"{reference_year}: {current_cap:,.2f} in {year}. " + "This usually means policyengine-us is missing the long-run NAWI/" + "payroll-cap extension." + ) + return current_cap + + def build_household_age_matrix(sim, n_ages=86): """ Build household age composition matrix from simulation. @@ -115,7 +145,14 @@ def get_pseudo_input_variables(sim): return pseudo_inputs -def create_household_year_h5(year, household_weights, base_dataset, output_dir): +def create_household_year_h5( + year, + household_weights, + base_dataset, + output_dir, + *, + reform=None, +): """ Create a year-specific .h5 file with calibrated household weights. @@ -127,13 +164,14 @@ def create_household_year_h5(year, household_weights, base_dataset, output_dir): household_weights: Calibrated household weights for this year base_dataset: Path to base dataset or in-memory Dataset instance output_dir: Directory to save the .h5 file + reform: Optional reform to apply when materializing year-specific values Returns: Path to the created .h5 file """ output_path = os.path.join(output_dir, f"{year}.h5") - sim = Microsimulation(dataset=base_dataset) + sim = Microsimulation(dataset=base_dataset, reform=reform) base_period = int(sim.default_calculation_period) df = sim.to_input_dataframe() diff --git a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py index 8450c5f9e..84339fbc3 100644 --- a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py +++ b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py @@ -2,6 +2,7 @@ import argparse from dataclasses import asdict, dataclass +from functools import lru_cache import json from pathlib import Path @@ -17,7 +18,11 @@ calibrate_entropy_bounded, densify_lp_solution, ) - from .projection_utils import aggregate_age_targets, build_age_bins + from .projection_utils import ( + aggregate_age_targets, + build_age_bins, + validate_projected_social_security_cap, + ) from .ssa_data import ( get_long_term_target_source, load_ssa_age_projections, @@ -32,7 +37,11 @@ calibrate_entropy_bounded, densify_lp_solution, ) - from projection_utils import aggregate_age_targets, build_age_bins + from projection_utils import ( + aggregate_age_targets, + build_age_bins, + validate_projected_social_security_cap, + ) from ssa_data import ( get_long_term_target_source, load_ssa_age_projections, @@ -102,6 +111,11 @@ class SyntheticCandidate: def payroll_total(self) -> float: return float(self.head_wages + self.spouse_wages) + def taxable_payroll_total(self, payroll_cap: float) -> float: + return float( + min(self.head_wages, payroll_cap) + min(self.spouse_wages, payroll_cap) + ) + @property def ss_total(self) -> float: return float(self.head_ss + self.spouse_ss) @@ -441,9 +455,11 @@ def build_tax_unit_summary( "social_security", period=period ).values, "payroll": ( - sim.calculate("employment_income_before_lsr", period=period).values + sim.calculate( + "taxable_earnings_for_social_security", period=period + ).values + sim.calculate( - "self_employment_income_before_lsr", period=period + "social_security_taxable_self_employment_income", period=period ).values ), "dividend_income": sim.calculate( @@ -618,6 +634,40 @@ def _scale_levels(levels: list[float], scale: float) -> list[float]: return [round(level * scale, 2) for level in levels] +@lru_cache(maxsize=None) +def load_policyengine_social_security_cap(year: int) -> float: + sim = Microsimulation(dataset=DEFAULT_DATASET) + return validate_projected_social_security_cap( + sim.tax_benefit_system.parameters, + year, + ) + + +def allocate_taxable_payroll_wages( + total_taxable_payroll: float, + payroll_split: tuple[float, float], + payroll_cap: float, + *, + has_spouse: bool, +) -> tuple[float, float]: + total_taxable_payroll = max(float(total_taxable_payroll), 0.0) + positive_earner_count = int(payroll_split[0] > 0) + int( + has_spouse and payroll_split[1] > 0 + ) + if not has_spouse or positive_earner_count <= 1: + if has_spouse and payroll_split[1] > payroll_split[0]: + return 0.0, min(total_taxable_payroll, payroll_cap) + return min(total_taxable_payroll, payroll_cap), 0.0 + + total_taxable_payroll = min(total_taxable_payroll, positive_earner_count * payroll_cap) + preferred_head = total_taxable_payroll * float(payroll_split[0]) + lower = max(0.0, total_taxable_payroll - payroll_cap) + upper = min(payroll_cap, total_taxable_payroll) + head_taxable = min(max(preferred_head, lower), upper) + spouse_taxable = total_taxable_payroll - head_taxable + return float(head_taxable), float(spouse_taxable) + + def build_quantile_pools( actual_summary: pd.DataFrame, *, @@ -705,6 +755,8 @@ def build_quantile_pools( def generate_synthetic_candidates( pools: dict[str, dict[str, list[float]]], + *, + payroll_cap: float, ) -> list[SyntheticCandidate]: candidates: list[SyntheticCandidate] = [] for template in TEMPLATES: @@ -736,13 +788,11 @@ def generate_synthetic_candidates( spouse_ss = ( scaled_ss_total * template.ss_split[1] ) - head_wages = ( - scaled_payroll_total - * template.payroll_split[0] - ) - spouse_wages = ( - scaled_payroll_total - * template.payroll_split[1] + head_wages, spouse_wages = allocate_taxable_payroll_wages( + scaled_payroll_total, + template.payroll_split, + payroll_cap, + has_spouse=spouse_age is not None, ) candidates.append( SyntheticCandidate( @@ -790,6 +840,80 @@ def age_bucket_vector(ages: list[int], age_bins: list[tuple[int, int]]) -> np.nd return vector +def build_role_composite_calibration_blueprint( + augmentation_report: dict[str, object], + *, + year: int, + age_bins: list[tuple[int, int]], + hh_id_to_idx: dict[int, int], + baseline_weights: np.ndarray, + base_weight_scale: float = 0.5, +) -> dict[str, object] | None: + """ + Build target-year calibration overrides for donor-composite clones. + + Donor-composite augmentation produces realized microdata rows that are close + to the synthetic target support but not numerically identical. At the + augmentation target year we can calibrate against the exact clone + blueprints, using the synthetic solution as priors, while still applying + the resulting household weights to the realized rows. + """ + target_year = augmentation_report.get("target_year") + clone_reports = augmentation_report.get("clone_household_reports") + if target_year is None or int(target_year) != int(year): + return None + if not clone_reports: + return None + + baseline_weights = np.asarray(baseline_weights, dtype=float) + if baseline_weights.ndim != 1: + raise ValueError("baseline_weights must be one-dimensional") + + prior_weights = np.maximum( + baseline_weights * float(base_weight_scale), + 1e-12, + ) + clone_total_prior_weight = max(float(baseline_weights.sum()), 1.0) + age_overrides: dict[int, np.ndarray] = {} + ss_overrides: dict[int, float] = {} + payroll_overrides: dict[int, float] = {} + applied_clone_households = 0 + + for clone_report in clone_reports: + household_id = int(clone_report["clone_household_id"]) + idx = hh_id_to_idx.get(household_id) + if idx is None: + continue + ages = [int(clone_report["target_head_age"])] + spouse_age = clone_report.get("target_spouse_age") + if spouse_age is not None: + ages.append(int(spouse_age)) + ages.extend(int(age) for age in clone_report.get("target_dependent_ages", [])) + age_overrides[idx] = age_bucket_vector(ages, age_bins) + ss_overrides[idx] = float(clone_report["target_ss_total"]) + payroll_overrides[idx] = float(clone_report["target_payroll_total"]) + prior_weights[idx] = ( + clone_total_prior_weight + * float(clone_report["per_clone_weight_share_pct"]) + / 100.0 + ) + applied_clone_households += 1 + + return { + "baseline_weights": np.maximum(prior_weights, 1e-12), + "age_overrides": age_overrides, + "ss_overrides": ss_overrides, + "payroll_overrides": payroll_overrides, + "summary": { + "mode": "target_year_role_composite_blueprint", + "target_year": int(target_year), + "clone_household_count": int(applied_clone_households), + "base_weight_scale": float(base_weight_scale), + "clone_total_prior_weight": float(clone_total_prior_weight), + }, + } + + def _ages_from_summary_row(row: pd.Series) -> list[int]: ages = [int(round(float(row["head_age"])))] if pd.notna(row.get("spouse_age")): @@ -1017,13 +1141,14 @@ def build_synthetic_constraint_problem( year: int, baseline_weights: np.ndarray | None = None, ) -> dict[str, object]: + payroll_cap = load_policyengine_social_security_cap(year) age_targets = load_ssa_age_projections(start_year=year, end_year=year) age_bins = build_age_bins(n_ages=age_targets.shape[0], bucket_size=5) aggregated_age_targets = aggregate_age_targets(age_targets, age_bins)[:, 0] X = np.vstack([age_bucket_vector(candidate.ages(), age_bins) for candidate in candidates]) ss_values = np.array([candidate.ss_total for candidate in candidates], dtype=float) payroll_values = np.array( - [candidate.payroll_total for candidate in candidates], + [candidate.taxable_payroll_total(payroll_cap) for candidate in candidates], dtype=float, ) if baseline_weights is None: @@ -1040,6 +1165,7 @@ def build_synthetic_constraint_problem( "ss_target": float(load_ssa_benefit_projections(year)), "payroll_values": payroll_values, "payroll_target": float(load_taxable_payroll_projections(year)), + "payroll_cap": float(payroll_cap), "baseline_weights": baseline_weights, } @@ -1339,6 +1465,7 @@ def build_role_donor_composites( composite_weights.append(per_candidate_weight) composite_records.append( { + "composite_idx": int(len(composite_candidates) - 1), "candidate_idx": int(target_row["candidate_idx"]), "archetype": target_candidate.archetype, "older_tax_unit_id": ( @@ -1383,7 +1510,7 @@ def build_role_donor_composites( "max_older_distance": float(max_older_distance), "max_worker_distance": float(max_worker_distance), "skipped_targets": skipped_targets, - "composite_records": composite_records[:100], + "composite_records": composite_records, "prior_summary": probe_summary, }, ) @@ -1834,22 +1961,34 @@ def _dependent_rows(df: pd.DataFrame | None) -> pd.DataFrame: elif len(older_adults) >= 2: spouse_row = older_adults.iloc[1].copy() if spouse_row is None: - return None, id_counters + fallback_spouse_pool = worker_adults if not worker_adults.empty else older_adults + if fallback_spouse_pool.empty: + return None, id_counters + spouse_row = fallback_spouse_pool.iloc[0].copy() selected_rows.append(spouse_row) if len(target_candidate.dependent_ages) > 0: - if len(worker_dependents) < len(target_candidate.dependent_ages): - return None, id_counters - selected_rows.extend( - [ - row.copy() - for _, row in worker_dependents.iloc[ - : len(target_candidate.dependent_ages) - ].iterrows() - ] - ) - - cloned = pd.DataFrame(selected_rows).copy() + dependent_rows = [row.copy() for _, row in worker_dependents.iterrows()] + if not dependent_rows: + fallback_source = None + if worker_donor_rows is not None and not worker_donor_rows.empty: + fallback_source = worker_donor_rows.sort_values(age_col, ascending=True).iloc[ + 0 + ].copy() + elif older_donor_rows is not None and not older_donor_rows.empty: + fallback_source = older_donor_rows.sort_values(age_col, ascending=True).iloc[ + 0 + ].copy() + if fallback_source is None: + return None, id_counters + dependent_rows = [fallback_source.copy()] + while len(dependent_rows) < len(target_candidate.dependent_ages): + dependent_rows.append(dependent_rows[-1].copy()) + selected_rows.extend(dependent_rows[: len(target_candidate.dependent_ages)]) + + # Reset duplicate donor indices so later row-specific retargeting only touches + # the intended clone row. + cloned = pd.DataFrame(selected_rows).reset_index(drop=True).copy() household_id = id_counters["household"] id_counters["household"] += 1 for entity_name, columns in ENTITY_ID_COLUMNS.items(): @@ -1989,7 +2128,10 @@ def build_donor_backed_augmented_input_dataframe( ss_scale=ss_scale, earnings_scale=earnings_scale, ) - candidates = generate_synthetic_candidates(pools) + candidates = generate_synthetic_candidates( + pools, + payroll_cap=load_policyengine_social_security_cap(target_year), + ) exact_weights, solve_info = solve_synthetic_support(candidates, year=target_year) exact_df = summarize_exact_candidates(candidates, exact_weights) target_df = exact_df[exact_df["synthetic_weight"] > 0].head(top_n_targets).copy() @@ -2121,15 +2263,48 @@ def build_role_composite_augmented_input_dataframe( ss_scale=ss_scale, earnings_scale=earnings_scale, ) - candidates = generate_synthetic_candidates(pools) + candidates = generate_synthetic_candidates( + pools, + payroll_cap=load_policyengine_social_security_cap(target_year), + ) exact_weights, solve_info = solve_synthetic_support(candidates, year=target_year) - exact_df = summarize_exact_candidates(candidates, exact_weights) - target_df = exact_df[exact_df["synthetic_weight"] > 0].head(top_n_targets).copy() scaled_actual = build_scaled_actual_summary( actual_summary, ss_scale=ss_scale, earnings_scale=earnings_scale, ) + ( + role_composite_candidates, + role_composite_prior, + role_composite_probe, + ) = build_role_donor_composites( + candidates, + exact_weights, + actual_summary, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + top_n_targets=top_n_targets, + older_donors_per_target=donors_per_target, + worker_donors_per_target=donors_per_target, + max_older_distance=max_older_distance, + max_worker_distance=max_worker_distance, + ) + role_composite_weights, role_composite_solve_info = solve_synthetic_support( + role_composite_candidates, + year=target_year, + baseline_weights=role_composite_prior, + ) + role_composite_df = summarize_exact_candidates( + role_composite_candidates, + role_composite_weights, + ) + selected_composite_df = role_composite_df[ + role_composite_df["synthetic_weight"] > 0 + ].copy() + composite_records_by_idx = { + int(record["composite_idx"]): record + for record in role_composite_probe["composite_records"] + } tax_unit_id_col = _period_column("person_tax_unit_id", base_year) id_counters = { @@ -2145,96 +2320,85 @@ def build_role_composite_augmented_input_dataframe( target_reports = [] skipped_targets = [] - for _, target_row in target_df.iterrows(): - target_candidate = candidates[int(target_row["candidate_idx"])] - older_matches = ( - match_older_role_donors( - target_candidate, - scaled_actual, - donors_per_target=donors_per_target, - ) - if target_candidate.ss_total > 0 - else pd.DataFrame([{}]) - ) - worker_matches = ( - match_worker_role_donors( - target_candidate, - scaled_actual, - donors_per_target=donors_per_target, + for _, target_row in selected_composite_df.iterrows(): + composite_idx = int(target_row["candidate_idx"]) + target_candidate = role_composite_candidates[composite_idx] + composite_record = composite_records_by_idx.get(composite_idx) + if composite_record is None: + skipped_targets.append( + { + "candidate_idx": composite_idx, + "archetype": target_candidate.archetype, + "weight_share_pct": float(target_row["weight_share_pct"]), + "reason": "missing_composite_record", + } ) - if target_candidate.payroll_total > 0 - else pd.DataFrame([{}]) - ) - if target_candidate.ss_total > 0: - older_matches = older_matches[ - older_matches["distance"] <= max_older_distance + continue + + older_tax_unit_id = composite_record.get("older_tax_unit_id") + worker_tax_unit_id = composite_record.get("worker_tax_unit_id") + older_row = None + worker_row = None + older_rows = None + worker_rows = None + if older_tax_unit_id is not None: + older_row = scaled_actual[ + scaled_actual["tax_unit_id"].eq(int(older_tax_unit_id)) + ].iloc[0] + older_rows = input_df[ + input_df[tax_unit_id_col] == int(older_tax_unit_id) ].copy() - if target_candidate.payroll_total > 0: - worker_matches = worker_matches[ - worker_matches["distance"] <= max_worker_distance + if worker_tax_unit_id is not None: + worker_row = scaled_actual[ + scaled_actual["tax_unit_id"].eq(int(worker_tax_unit_id)) + ].iloc[0] + worker_rows = input_df[ + input_df[tax_unit_id_col] == int(worker_tax_unit_id) ].copy() - if older_matches.empty or worker_matches.empty: + + clone_df, id_counters = _compose_role_donor_rows_to_target( + older_rows, + worker_rows, + base_year=base_year, + target_candidate=target_candidate, + ss_scale=ss_scale, + earnings_scale=earnings_scale, + id_counters=id_counters, + clone_weight_scale=clone_weight_scale, + clone_weight_divisor=1, + ) + if clone_df is None: skipped_targets.append( { - "candidate_idx": int(target_row["candidate_idx"]), + "candidate_idx": composite_idx, "archetype": target_candidate.archetype, "weight_share_pct": float(target_row["weight_share_pct"]), - "older_match_count": int(len(older_matches)), - "worker_match_count": int(len(worker_matches)), + "reason": "clone_build_failed", } ) continue - - successful_clone_count = 0 - combination_count = len(older_matches) * len(worker_matches) - for _, older_row in older_matches.iterrows(): - older_rows = ( - input_df[input_df[tax_unit_id_col] == int(older_row["tax_unit_id"])].copy() - if target_candidate.ss_total > 0 - else None + clone_frames.append(clone_df) + clone_household_reports.append( + _clone_report_record( + clone_df=clone_df, + base_year=base_year, + target_candidate=target_candidate, + candidate_idx=composite_idx, + target_weight_share_pct=float(target_row["weight_share_pct"]), + clone_weight_scale=clone_weight_scale, + combination_count=1, + older_donor_row=older_row, + worker_donor_row=worker_row, ) - for _, worker_row in worker_matches.iterrows(): - worker_rows = ( - input_df[input_df[tax_unit_id_col] == int(worker_row["tax_unit_id"])].copy() - if target_candidate.payroll_total > 0 - else None - ) - clone_df, id_counters = _compose_role_donor_rows_to_target( - older_rows, - worker_rows, - base_year=base_year, - target_candidate=target_candidate, - ss_scale=ss_scale, - earnings_scale=earnings_scale, - id_counters=id_counters, - clone_weight_scale=clone_weight_scale, - clone_weight_divisor=combination_count, - ) - if clone_df is None: - continue - clone_frames.append(clone_df) - clone_household_reports.append( - _clone_report_record( - clone_df=clone_df, - base_year=base_year, - target_candidate=target_candidate, - candidate_idx=int(target_row["candidate_idx"]), - target_weight_share_pct=float(target_row["weight_share_pct"]), - clone_weight_scale=clone_weight_scale, - combination_count=combination_count, - older_donor_row=older_row if target_candidate.ss_total > 0 else None, - worker_donor_row=worker_row if target_candidate.payroll_total > 0 else None, - ) - ) - successful_clone_count += 1 + ) target_reports.append( { - "candidate_idx": int(target_row["candidate_idx"]), + "candidate_idx": composite_idx, "archetype": target_candidate.archetype, "weight_share_pct": float(target_row["weight_share_pct"]), - "older_match_count": int(len(older_matches)), - "worker_match_count": int(len(worker_matches)), - "successful_clone_count": int(successful_clone_count), + "older_match_count": int(older_tax_unit_id is not None), + "worker_match_count": int(worker_tax_unit_id is not None), + "successful_clone_count": 1, } ) @@ -2255,6 +2419,8 @@ def build_role_composite_augmented_input_dataframe( "target_year": int(target_year), "target_source": get_long_term_target_source(), "solve_info": solve_info, + "role_composite_solve_info": role_composite_solve_info, + "selection_strategy": "role_composite_positive_support", "top_n_targets": int(top_n_targets), "donors_per_target": int(donors_per_target), "max_older_distance": float(max_older_distance), @@ -2268,6 +2434,8 @@ def build_role_composite_augmented_input_dataframe( ), "base_person_count": int(len(input_df)), "augmented_person_count": int(len(augmented_df)), + "role_composite_candidate_count": int(len(role_composite_candidates)), + "selected_role_composite_count": int(len(selected_composite_df)), "clone_household_count": int(len(clone_household_reports)), "clone_household_reports": clone_household_reports, "target_reports": target_reports, @@ -2818,7 +2986,10 @@ def main() -> int: ss_scale=ss_scale, earnings_scale=earnings_scale, ) - candidates = generate_synthetic_candidates(pools) + candidates = generate_synthetic_candidates( + pools, + payroll_cap=load_policyengine_social_security_cap(args.year), + ) weights, solve_info = solve_synthetic_support(candidates, year=args.year) solution_summary = summarize_solution(candidates, weights, actual_summary) donor_probe = summarize_donor_probe( diff --git a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py index ba54bc839..f90b51f68 100644 --- a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py +++ b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py @@ -3,18 +3,22 @@ Usage: - python run_household_projection.py [START_YEAR] [END_YEAR] [--profile PROFILE] [--target-source SOURCE] [--output-dir DIR] [--save-h5] [--allow-validation-failures] + python run_household_projection.py [START_YEAR] [END_YEAR] [--profile PROFILE] [--target-source SOURCE] [--tax-assumption ASSUMPTION] [--output-dir DIR] [--save-h5] [--allow-validation-failures] python run_household_projection.py [START_YEAR] [END_YEAR] [--profile PROFILE] [--target-source SOURCE] [--support-augmentation-profile donor-backed-synthetic-v1] [--support-augmentation-target-year YEAR] + python run_household_projection.py [START_YEAR] [END_YEAR] [--profile PROFILE] [--target-source SOURCE] [--support-augmentation-profile donor-backed-composite-v1] [--support-augmentation-target-year YEAR] [--support-augmentation-align-to-run-year] [--support-augmentation-blueprint-base-weight-scale SCALE] python run_household_projection.py [START_YEAR] [END_YEAR] [--greg] [--use-ss] [--use-payroll] [--use-h6-reform] [--use-tob] [--save-h5] START_YEAR: Optional starting year (default: 2025) END_YEAR: Optional ending year (default: 2035) --profile: Named calibration contract (recommended) --target-source: Named long-term target source package + --tax-assumption: Long-run federal tax assumption (`trustees-core-thresholds-v1` by default) --output-dir: Output directory for generated H5 files and metadata --allow-validation-failures: Record validation issues in metadata and continue instead of aborting the run --support-augmentation-profile: Experimental late-year support expansion profile --support-augmentation-target-year: Year whose extreme support is used to build the supplement + --support-augmentation-align-to-run-year: Rebuild the augmentation support for each run year instead of reusing one target-year support + --support-augmentation-blueprint-base-weight-scale: Prior scaling applied to original households when target-year donor-composite blueprint calibration is active --greg: Use GREG calibration instead of IPF (optional) --use-ss: Include Social Security benefit totals as calibration target (requires --greg) --use-payroll: Include taxable payroll totals as calibration target (requires --greg) @@ -26,6 +30,8 @@ python run_household_projection.py 2045 2045 --profile ss --target-source trustees_2025_current_law --save-h5 python run_household_projection.py 2025 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --save-h5 python run_household_projection.py 2075 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --support-augmentation-profile donor-backed-synthetic-v1 --support-augmentation-target-year 2100 --allow-validation-failures + python run_household_projection.py 2100 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --support-augmentation-profile donor-backed-composite-v1 --support-augmentation-target-year 2100 --support-augmentation-blueprint-base-weight-scale 0.5 --allow-validation-failures --save-h5 + python run_household_projection.py 2075 2100 --profile ss-payroll-tob --target-source trustees_2025_current_law --support-augmentation-profile donor-backed-composite-v1 --support-augmentation-align-to-run-year --support-augmentation-blueprint-base-weight-scale 0.5 --allow-validation-failures """ import sys @@ -64,8 +70,15 @@ build_age_bins, build_household_age_matrix, create_household_year_h5, + validate_projected_social_security_cap, +) +from tax_assumptions import ( + TRUSTEES_CORE_THRESHOLD_ASSUMPTION, + create_wage_indexed_core_thresholds_reform, + get_long_run_tax_assumption_metadata, ) from prototype_synthetic_2100_support import ( + build_role_composite_calibration_blueprint, build_donor_backed_augmented_dataset, build_role_composite_augmented_dataset, ) @@ -261,6 +274,10 @@ def create_h6_reform(): "donor-backed-synthetic-v1", "donor-backed-composite-v1", } +SUPPORTED_TAX_ASSUMPTIONS = { + "current-law-literal", + TRUSTEES_CORE_THRESHOLD_ASSUMPTION["name"], +} PROFILE_NAME = None @@ -307,6 +324,12 @@ def create_h6_reform(): SUPPORT_AUGMENTATION_TARGET_YEAR = int(sys.argv[target_year_index + 1]) del sys.argv[target_year_index : target_year_index + 2] +SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR = ( + "--support-augmentation-align-to-run-year" in sys.argv +) +if SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR: + sys.argv.remove("--support-augmentation-align-to-run-year") + SUPPORT_AUGMENTATION_START_YEAR = 2075 if "--support-augmentation-start-year" in sys.argv: start_year_index = sys.argv.index("--support-augmentation-start-year") @@ -361,6 +384,33 @@ def create_h6_reform(): ) del sys.argv[weight_scale_index : weight_scale_index + 2] +SUPPORT_AUGMENTATION_BLUEPRINT_BASE_WEIGHT_SCALE = 0.5 +if "--support-augmentation-blueprint-base-weight-scale" in sys.argv: + blueprint_scale_index = sys.argv.index( + "--support-augmentation-blueprint-base-weight-scale" + ) + if blueprint_scale_index + 1 >= len(sys.argv): + raise ValueError( + "--support-augmentation-blueprint-base-weight-scale requires a float" + ) + SUPPORT_AUGMENTATION_BLUEPRINT_BASE_WEIGHT_SCALE = float( + sys.argv[blueprint_scale_index + 1] + ) + del sys.argv[blueprint_scale_index : blueprint_scale_index + 2] + +TAX_ASSUMPTION = TRUSTEES_CORE_THRESHOLD_ASSUMPTION["name"] +if "--tax-assumption" in sys.argv: + tax_assumption_index = sys.argv.index("--tax-assumption") + if tax_assumption_index + 1 >= len(sys.argv): + raise ValueError("--tax-assumption requires a value") + TAX_ASSUMPTION = sys.argv[tax_assumption_index + 1] + del sys.argv[tax_assumption_index : tax_assumption_index + 2] +if TAX_ASSUMPTION not in SUPPORTED_TAX_ASSUMPTIONS: + raise ValueError( + "Unsupported --tax-assumption: " + f"{TAX_ASSUMPTION}. Valid values: {sorted(SUPPORTED_TAX_ASSUMPTIONS)}" + ) + ALLOW_VALIDATION_FAILURES = "--allow-validation-failures" in sys.argv if ALLOW_VALIDATION_FAILURES: sys.argv.remove("--allow-validation-failures") @@ -444,8 +494,42 @@ def create_h6_reform(): TARGET_SOURCE = get_long_term_target_source() TARGET_SOURCE_METADATA = describe_long_term_target_source(TARGET_SOURCE) + +def _compose_reforms(*reforms): + reforms = tuple(reform for reform in reforms if reform is not None) + if not reforms: + return None + if len(reforms) == 1: + return reforms[0] + return reforms + + +if TAX_ASSUMPTION == "current-law-literal": + ACTIVE_LONG_RUN_TAX_REFORM = None + LONG_RUN_TAX_ASSUMPTION_METADATA = { + "name": "current-law-literal", + "description": ( + "Use the baseline PolicyEngine federal tax parameter uprating " + "without long-run Trustees-style wage-index overrides." + ), + "source": "PolicyEngine baseline", + "start_year": None, + "end_year": int(END_YEAR), + } +else: + ACTIVE_LONG_RUN_TAX_REFORM = create_wage_indexed_core_thresholds_reform( + start_year=TRUSTEES_CORE_THRESHOLD_ASSUMPTION["start_year"], + end_year=END_YEAR, + ) + LONG_RUN_TAX_ASSUMPTION_METADATA = get_long_run_tax_assumption_metadata( + TAX_ASSUMPTION, + end_year=END_YEAR, + ) + BASE_DATASET = BASE_DATASET_PATH SUPPORT_AUGMENTATION_METADATA = None +SUPPORT_AUGMENTATION_REPORT = None +MANIFEST_SUPPORT_AUGMENTATION_METADATA = None CALIBRATION_METHOD = PROFILE.calibration_method USE_GREG = CALIBRATION_METHOD == "greg" @@ -483,12 +567,20 @@ def create_h6_reform(): print(f" Calibration profile: {PROFILE.name}") print(f" Profile description: {PROFILE.description}") print(f" Target source: {TARGET_SOURCE}") +print(f" Long-run tax assumption: {LONG_RUN_TAX_ASSUMPTION_METADATA['name']}") print(f" Calibration method: {CALIBRATION_METHOD.upper()}") if SUPPORT_AUGMENTATION_PROFILE: print(f" Support augmentation: {SUPPORT_AUGMENTATION_PROFILE}") + if SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR: + print(" Support augmentation target year: each run year") + else: + print( + " Support augmentation target year: " + f"{SUPPORT_AUGMENTATION_TARGET_YEAR}" + ) print( - " Support augmentation target year: " - f"{SUPPORT_AUGMENTATION_TARGET_YEAR}" + " Support augmentation blueprint base-weight scale: " + f"{SUPPORT_AUGMENTATION_BLUEPRINT_BASE_WEIGHT_SCALE}" ) if USE_SS: print(f" Including Social Security benefits constraint: Yes") @@ -509,75 +601,61 @@ def create_h6_reform(): est_time = (END_YEAR - START_YEAR + 1) * (3 if SAVE_H5 else 2) print(f" Estimated time: ~{est_time:.0f} minutes") -# ========================================================================= -# STEP 1: LOAD SSA DEMOGRAPHIC PROJECTIONS -# ========================================================================= -print("\n" + "=" * 70) -print("STEP 1: DEMOGRAPHIC PROJECTIONS") -print("=" * 70) -target_matrix = load_ssa_age_projections(start_year=START_YEAR, end_year=END_YEAR) -n_years = target_matrix.shape[1] -n_ages = target_matrix.shape[0] +def _build_support_augmentation( + target_year: int, + *, + report_filename: str | None = None, +): + if SUPPORT_AUGMENTATION_PROFILE is None: + return BASE_DATASET_PATH, None, None, None -print(f"\nLoaded SSA projections: {n_ages} ages x {n_years} years") -print(f"\nPopulation projections:") - -display_years = [ - y - for y in [START_YEAR, 2030, 2040, 2050, 2060, 2070, 2080, 2090, 2100] - if START_YEAR <= y <= END_YEAR -] -if END_YEAR not in display_years: - display_years.append(END_YEAR) - -for y in display_years: - idx = y - START_YEAR - if idx < n_years: - pop = target_matrix[:, idx].sum() - print(f" {y}: {pop / 1e6:6.1f}M") - -if SUPPORT_AUGMENTATION_PROFILE in { - "donor-backed-synthetic-v1", - "donor-backed-composite-v1", -}: - print("\n" + "=" * 70) - print("STEP 1B: BUILD DONOR-BACKED LATE-YEAR SUPPORT") - print("=" * 70) if SUPPORT_AUGMENTATION_PROFILE == "donor-backed-synthetic-v1": - BASE_DATASET, augmentation_report = build_donor_backed_augmented_dataset( + augmented_dataset, augmentation_report = build_donor_backed_augmented_dataset( base_dataset=BASE_DATASET_PATH, base_year=BASE_YEAR, - target_year=SUPPORT_AUGMENTATION_TARGET_YEAR, + target_year=target_year, top_n_targets=SUPPORT_AUGMENTATION_TOP_N_TARGETS, donors_per_target=SUPPORT_AUGMENTATION_DONORS_PER_TARGET, max_distance_for_clone=SUPPORT_AUGMENTATION_MAX_DISTANCE, clone_weight_scale=SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE, ) else: - BASE_DATASET, augmentation_report = build_role_composite_augmented_dataset( + augmented_dataset, augmentation_report = build_role_composite_augmented_dataset( base_dataset=BASE_DATASET_PATH, base_year=BASE_YEAR, - target_year=SUPPORT_AUGMENTATION_TARGET_YEAR, + target_year=target_year, top_n_targets=SUPPORT_AUGMENTATION_TOP_N_TARGETS, donors_per_target=SUPPORT_AUGMENTATION_DONORS_PER_TARGET, max_older_distance=SUPPORT_AUGMENTATION_MAX_DISTANCE, max_worker_distance=SUPPORT_AUGMENTATION_MAX_DISTANCE, clone_weight_scale=SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE, ) - support_augmentation_report_path = write_support_augmentation_report( + + report_path = write_support_augmentation_report( OUTPUT_DIR, augmentation_report, + filename=( + report_filename + if report_filename is not None + else "support_augmentation_report.json" + ), ) - SUPPORT_AUGMENTATION_METADATA = { + year_metadata = { "name": SUPPORT_AUGMENTATION_PROFILE, "activation_start_year": SUPPORT_AUGMENTATION_START_YEAR, - "target_year": SUPPORT_AUGMENTATION_TARGET_YEAR, + "target_year": int(target_year), + "target_year_strategy": ( + "run_year" if SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR else "fixed" + ), "top_n_targets": SUPPORT_AUGMENTATION_TOP_N_TARGETS, "donors_per_target": SUPPORT_AUGMENTATION_DONORS_PER_TARGET, "max_distance_for_clone": SUPPORT_AUGMENTATION_MAX_DISTANCE, "clone_weight_scale": SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE, - "report_file": support_augmentation_report_path.name, + "blueprint_base_weight_scale": ( + SUPPORT_AUGMENTATION_BLUEPRINT_BASE_WEIGHT_SCALE + ), + "report_file": report_path.name, "report_summary": { "base_household_count": augmentation_report["base_household_count"], "augmented_household_count": augmentation_report[ @@ -587,16 +665,42 @@ def create_h6_reform(): "augmented_person_count": augmentation_report[ "augmented_person_count" ], - "clone_household_count": augmentation_report.get("clone_household_count", 0), + "clone_household_count": augmentation_report.get( + "clone_household_count", 0 + ), "successful_target_count": sum( report["successful_clone_count"] > 0 for report in augmentation_report["target_reports"] ), - "skipped_target_count": len( - augmentation_report["skipped_targets"] - ), + "skipped_target_count": len(augmentation_report["skipped_targets"]), }, } + manifest_metadata = { + "name": SUPPORT_AUGMENTATION_PROFILE, + "activation_start_year": SUPPORT_AUGMENTATION_START_YEAR, + "target_year": ( + int(target_year) if not SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR else None + ), + "target_year_strategy": ( + "run_year" if SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR else "fixed" + ), + "top_n_targets": SUPPORT_AUGMENTATION_TOP_N_TARGETS, + "donors_per_target": SUPPORT_AUGMENTATION_DONORS_PER_TARGET, + "max_distance_for_clone": SUPPORT_AUGMENTATION_MAX_DISTANCE, + "clone_weight_scale": SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE, + "blueprint_base_weight_scale": ( + SUPPORT_AUGMENTATION_BLUEPRINT_BASE_WEIGHT_SCALE + ), + "report_file": ( + None + if SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR + else report_path.name + ), + } + return augmented_dataset, augmentation_report, year_metadata, manifest_metadata + + +def _print_support_augmentation_summary(augmentation_report: dict) -> None: print( " Base households -> augmented households: " f"{augmentation_report['base_household_count']:,} -> " @@ -616,6 +720,86 @@ def create_h6_reform(): f"{len(augmentation_report['skipped_targets'])}" ) +# ========================================================================= +# STEP 1: LOAD SSA DEMOGRAPHIC PROJECTIONS +# ========================================================================= +print("\n" + "=" * 70) +print("STEP 1: DEMOGRAPHIC PROJECTIONS") +print("=" * 70) + +target_matrix = load_ssa_age_projections(start_year=START_YEAR, end_year=END_YEAR) +n_years = target_matrix.shape[1] +n_ages = target_matrix.shape[0] + +print(f"\nLoaded SSA projections: {n_ages} ages x {n_years} years") +print(f"\nPopulation projections:") + +display_years = [ + y + for y in [START_YEAR, 2030, 2040, 2050, 2060, 2070, 2080, 2090, 2100] + if START_YEAR <= y <= END_YEAR +] +if END_YEAR not in display_years: + display_years.append(END_YEAR) + +for y in display_years: + idx = y - START_YEAR + if idx < n_years: + pop = target_matrix[:, idx].sum() + print(f" {y}: {pop / 1e6:6.1f}M") + +augmentation_cache: dict[int, tuple[object, dict, dict, dict]] = {} +X = None +hh_id_to_idx = None +n_households = None +household_ids_unique = None +aggregated_age_cache: dict[int, tuple[np.ndarray, np.ndarray]] = {} + +if SUPPORT_AUGMENTATION_PROFILE in { + "donor-backed-synthetic-v1", + "donor-backed-composite-v1", +}: + print("\n" + "=" * 70) + print("STEP 1B: BUILD DONOR-BACKED LATE-YEAR SUPPORT") + print("=" * 70) + if SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR: + print( + " Dynamic mode: donor-backed support will be rebuilt separately for " + "each run year." + ) + else: + ( + BASE_DATASET, + SUPPORT_AUGMENTATION_REPORT, + SUPPORT_AUGMENTATION_METADATA, + MANIFEST_SUPPORT_AUGMENTATION_METADATA, + ) = _build_support_augmentation( + SUPPORT_AUGMENTATION_TARGET_YEAR, + ) + augmentation_cache[SUPPORT_AUGMENTATION_TARGET_YEAR] = ( + BASE_DATASET, + SUPPORT_AUGMENTATION_REPORT, + SUPPORT_AUGMENTATION_METADATA, + MANIFEST_SUPPORT_AUGMENTATION_METADATA, + ) + _print_support_augmentation_summary(SUPPORT_AUGMENTATION_REPORT) + +if SUPPORT_AUGMENTATION_PROFILE is not None and SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR: + MANIFEST_SUPPORT_AUGMENTATION_METADATA = { + "name": SUPPORT_AUGMENTATION_PROFILE, + "activation_start_year": SUPPORT_AUGMENTATION_START_YEAR, + "target_year": None, + "target_year_strategy": "run_year", + "top_n_targets": SUPPORT_AUGMENTATION_TOP_N_TARGETS, + "donors_per_target": SUPPORT_AUGMENTATION_DONORS_PER_TARGET, + "max_distance_for_clone": SUPPORT_AUGMENTATION_MAX_DISTANCE, + "clone_weight_scale": SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE, + "blueprint_base_weight_scale": ( + SUPPORT_AUGMENTATION_BLUEPRINT_BASE_WEIGHT_SCALE + ), + "report_file": None, + } + # ========================================================================= # STEP 2: BUILD HOUSEHOLD AGE MATRIX # ========================================================================= @@ -623,10 +807,12 @@ def create_h6_reform(): print("STEP 2: BUILDING HOUSEHOLD AGE COMPOSITION") print("=" * 70) +if SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR: + print("\nDynamic augmentation enabled; base support will be used before the activation year and rebuilt per-year after that.") + sim = Microsimulation(dataset=BASE_DATASET) X, household_ids_unique, hh_id_to_idx = build_household_age_matrix(sim, n_ages) n_households = len(household_ids_unique) -aggregated_age_cache: dict[int, tuple[np.ndarray, np.ndarray]] = {} print(f"\nLoaded {n_households:,} households") print(f"Household age matrix shape: {X.shape}") @@ -650,8 +836,6 @@ def create_h6_reform(): total_income_tax = np.zeros(n_years) total_income_tax_baseline = np.zeros(n_years) total_population = np.zeros(n_years) -weights_matrix = np.zeros((n_households, n_years)) -baseline_weights_matrix = np.zeros((n_households, n_years)) process = psutil.Process() print(f"\nInitial memory usage: {process.memory_info().rss / 1024**3:.2f} GB") @@ -661,8 +845,45 @@ def create_h6_reform(): for year_idx in range(n_years): year = START_YEAR + year_idx - - sim = Microsimulation(dataset=BASE_DATASET) + current_dataset = BASE_DATASET + current_support_augmentation_report = SUPPORT_AUGMENTATION_REPORT + current_support_augmentation_metadata = SUPPORT_AUGMENTATION_METADATA + current_X = X + current_hh_id_to_idx = hh_id_to_idx + + if ( + SUPPORT_AUGMENTATION_PROFILE is not None + and SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR + and year >= SUPPORT_AUGMENTATION_START_YEAR + ): + cached = augmentation_cache.get(year) + if cached is None: + cached = _build_support_augmentation( + year, + report_filename=f"support_augmentation_report_{year}.json", + ) + augmentation_cache[year] = cached + ( + current_dataset, + current_support_augmentation_report, + current_support_augmentation_metadata, + _, + ) = cached + if year in display_years: + print( + f" [DEBUG {year}] Rebuilt support augmentation for run year {year}" + ) + _print_support_augmentation_summary(current_support_augmentation_report) + sim = Microsimulation( + dataset=current_dataset, + reform=ACTIVE_LONG_RUN_TAX_REFORM, + ) + current_X, _, current_hh_id_to_idx = build_household_age_matrix(sim, n_ages) + else: + sim = Microsimulation( + dataset=current_dataset, + reform=ACTIVE_LONG_RUN_TAX_REFORM, + ) income_tax_hh = sim.calculate("income_tax", period=year, map_to="household") income_tax_baseline_total = income_tax_hh.sum() @@ -672,7 +893,11 @@ def create_h6_reform(): baseline_weights = household_microseries.weights.values household_ids_hh = household_microseries.values - assert len(household_ids_hh) == n_households + if not ( + SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR + and current_support_augmentation_report is not None + ): + assert len(household_ids_hh) == n_households ss_values = None ss_target = None @@ -689,6 +914,10 @@ def create_h6_reform(): payroll_values = None payroll_target = None if USE_PAYROLL: + payroll_cap = validate_projected_social_security_cap( + sim.tax_benefit_system.parameters, + year, + ) # SSA taxable payroll = W-2 wages capped at wage base + SE income within remaining cap room taxable_wages_hh = sim.calculate( "taxable_earnings_for_social_security", @@ -704,6 +933,9 @@ def create_h6_reform(): payroll_target = load_taxable_payroll_projections(year) if year in display_years: payroll_baseline = np.sum(payroll_values * baseline_weights) + print( + f" [DEBUG {year}] Payroll cap: ${payroll_cap:,.0f}" + ) print( f" [DEBUG {year}] Payroll baseline: ${payroll_baseline / 1e9:.1f}B, target: ${payroll_target / 1e9:.1f}B" ) @@ -722,7 +954,10 @@ def create_h6_reform(): else: # Create and apply H6 reform h6_reform = create_h6_reform() - reform_sim = Microsimulation(dataset=BASE_DATASET, reform=h6_reform) + reform_sim = Microsimulation( + dataset=current_dataset, + reform=_compose_reforms(ACTIVE_LONG_RUN_TAX_REFORM, h6_reform), + ) # Calculate reform income tax income_tax_reform_hh = reform_sim.calculate( @@ -785,28 +1020,80 @@ def create_h6_reform(): else None ) if age_bucket_size and age_bucket_size > 1: - if age_bucket_size not in aggregated_age_cache: + if SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR and current_support_augmentation_report is not None: age_bins = build_age_bins(n_ages=n_ages, bucket_size=age_bucket_size) - aggregated_age_cache[age_bucket_size] = ( - aggregate_household_age_matrix(X, age_bins), - aggregate_age_targets(target_matrix, age_bins), - ) - X_current, aggregated_target_matrix = aggregated_age_cache[age_bucket_size] + X_current = aggregate_household_age_matrix(current_X, age_bins) + aggregated_target_matrix = aggregate_age_targets(target_matrix, age_bins) + else: + if age_bucket_size not in aggregated_age_cache: + age_bins = build_age_bins(n_ages=n_ages, bucket_size=age_bucket_size) + aggregated_age_cache[age_bucket_size] = ( + aggregate_household_age_matrix(current_X, age_bins), + aggregate_age_targets(target_matrix, age_bins), + ) + X_current, aggregated_target_matrix = aggregated_age_cache[age_bucket_size] y_target = aggregated_target_matrix[:, year_idx] else: - X_current = X + X_current = current_X y_target = target_matrix[:, year_idx] age_bucket_size = 1 + X_actual_current = X_current + ss_values_actual = None if ss_values is None else np.asarray(ss_values, dtype=float) + payroll_values_actual = ( + None if payroll_values is None else np.asarray(payroll_values, dtype=float) + ) + calibration_baseline_weights = baseline_weights + X_calibration = X_current.copy() + ss_values_calibration = ( + None if ss_values_actual is None else ss_values_actual.copy() + ) + payroll_values_calibration = ( + None if payroll_values_actual is None else payroll_values_actual.copy() + ) + blueprint_summary = None + if ( + SUPPORT_AUGMENTATION_PROFILE == "donor-backed-composite-v1" + and current_support_augmentation_report is not None + and year >= SUPPORT_AUGMENTATION_START_YEAR + ): + calibration_blueprint = build_role_composite_calibration_blueprint( + current_support_augmentation_report, + year=year, + age_bins=build_age_bins(n_ages=n_ages, bucket_size=age_bucket_size), + hh_id_to_idx=current_hh_id_to_idx, + baseline_weights=baseline_weights, + base_weight_scale=SUPPORT_AUGMENTATION_BLUEPRINT_BASE_WEIGHT_SCALE, + ) + if calibration_blueprint is not None: + calibration_baseline_weights = calibration_blueprint["baseline_weights"] + for idx, age_vector in calibration_blueprint["age_overrides"].items(): + X_calibration[idx] = age_vector + if ss_values_calibration is not None: + for idx, target_value in calibration_blueprint["ss_overrides"].items(): + ss_values_calibration[idx] = target_value + if payroll_values_calibration is not None: + for idx, target_value in calibration_blueprint[ + "payroll_overrides" + ].items(): + payroll_values_calibration[idx] = target_value + blueprint_summary = calibration_blueprint["summary"] + if year in display_years: + print( + f" [DEBUG {year}] Using support blueprint for " + f"{blueprint_summary['clone_household_count']} clone households " + f"(base-weight scale {blueprint_summary['base_weight_scale']:.3f})" + ) + w_new, iterations, calibration_event = calibrate_weights( - X=X_current, + X=X_calibration, y_target=y_target, - baseline_weights=baseline_weights, + baseline_weights=calibration_baseline_weights, method=CALIBRATION_METHOD, calibrator=calibrator, - ss_values=ss_values, + ss_values=ss_values_calibration, ss_target=ss_target, - payroll_values=payroll_values, + payroll_values=payroll_values_calibration, payroll_target=payroll_target, h6_income_values=h6_income_values, h6_revenue_target=h6_revenue_target, @@ -828,14 +1115,14 @@ def create_h6_reform(): ) calibration_audit = build_calibration_audit( - X=X_current, + X=X_actual_current, y_target=y_target, weights=w_new, - baseline_weights=baseline_weights, + baseline_weights=calibration_baseline_weights, calibration_event=calibration_event, - ss_values=ss_values, + ss_values=ss_values_actual, ss_target=ss_target, - payroll_values=payroll_values, + payroll_values=payroll_values_actual, payroll_target=payroll_target, h6_income_values=h6_income_values, h6_revenue_target=h6_revenue_target, @@ -844,6 +1131,8 @@ def create_h6_reform(): hi_tob_values=hi_tob_values if USE_TOB else None, hi_tob_target=hi_tob_target if USE_TOB else None, ) + if blueprint_summary is not None: + calibration_audit["support_blueprint"] = blueprint_summary if BENCHMARK_TOB and oasdi_tob_values is not None and hi_tob_values is not None: calibration_audit["benchmarks"] = { "oasdi_tob": { @@ -950,14 +1239,18 @@ def create_h6_reform(): f"{hi_stats['pct_error']:.3f}%)" ) - weights_matrix[:, year_idx] = w_new - baseline_weights_matrix[:, year_idx] = baseline_weights total_income_tax[year_idx] = np.sum(income_tax_values * w_new) total_income_tax_baseline[year_idx] = income_tax_baseline_total total_population[year_idx] = np.sum(y_target) if SAVE_H5: - h5_path = create_household_year_h5(year, w_new, BASE_DATASET, OUTPUT_DIR) + h5_path = create_household_year_h5( + year, + w_new, + current_dataset, + OUTPUT_DIR, + reform=ACTIVE_LONG_RUN_TAX_REFORM, + ) metadata_path = write_year_metadata( h5_path, year=year, @@ -965,7 +1258,8 @@ def create_h6_reform(): profile=PROFILE.to_dict(), calibration_audit=calibration_audit, target_source=TARGET_SOURCE_METADATA, - support_augmentation=SUPPORT_AUGMENTATION_METADATA, + tax_assumption=LONG_RUN_TAX_ASSUMPTION_METADATA, + support_augmentation=current_support_augmentation_metadata, ) update_dataset_manifest( OUTPUT_DIR, @@ -976,7 +1270,8 @@ def create_h6_reform(): profile=PROFILE.to_dict(), calibration_audit=calibration_audit, target_source=TARGET_SOURCE_METADATA, - support_augmentation=SUPPORT_AUGMENTATION_METADATA, + tax_assumption=LONG_RUN_TAX_ASSUMPTION_METADATA, + support_augmentation=MANIFEST_SUPPORT_AUGMENTATION_METADATA, ) if year in display_years: print(f" Saved {year}.h5 and metadata") diff --git a/policyengine_us_data/datasets/cps/long_term/tax_assumptions.py b/policyengine_us_data/datasets/cps/long_term/tax_assumptions.py new file mode 100644 index 000000000..e93f4580e --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/tax_assumptions.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import math +from typing import Any + + +TRUSTEES_CORE_THRESHOLD_ASSUMPTION = { + "name": "trustees-core-thresholds-v1", + "description": ( + "Best-public Trustees tax-side approximation: keep Social Security " + "benefit-tax thresholds fixed, but wage-index core ordinary federal " + "tax thresholds after 2034." + ), + "source": "SSA 2025 Trustees Report V.C.7", + "start_year": 2035, + "parameter_groups": [ + "ordinary_income_brackets", + "standard_deduction", + "aged_blind_standard_deduction", + "capital_gains_thresholds", + "amt_thresholds", + ], +} + + +def round_amount(amount: float, rounding: dict | None) -> float: + if not rounding: + return amount + + interval = float(rounding["interval"]) + rounding_type = rounding["type"] + + if rounding_type == "downwards": + return math.floor(amount / interval) * interval + if rounding_type == "nearest": + return math.floor(amount / interval + 0.5) * interval + + raise ValueError(f"Unsupported rounding type: {rounding_type}") + + +def _uprating_parameter_name(parameter) -> str | None: + metadata = getattr(parameter, "metadata", {}) + uprating = metadata.get("uprating") + if isinstance(uprating, dict): + return uprating.get("parameter") + return uprating + + +def iter_updatable_parameters( + root, + *, + uprating_parameter: str | None = None, +) -> list: + candidates = [root] + if hasattr(root, "get_descendants"): + candidates.extend(root.get_descendants()) + + result = [] + for candidate in candidates: + if candidate.__class__.__name__ != "Parameter": + continue + uprating_name = _uprating_parameter_name(candidate) + if uprating_name is None: + continue + if uprating_parameter is not None and uprating_name != uprating_parameter: + continue + result.append(candidate) + return result + + +def apply_wage_growth_to_parameter( + parameter, + *, + nawi, + start_year: int, + end_year: int, +) -> None: + metadata = getattr(parameter, "metadata", {}) + uprating = metadata.get("uprating") + rounding = uprating.get("rounding") if isinstance(uprating, dict) else None + + for year in range(start_year, end_year + 1): + previous_value = float(parameter(f"{year - 1}-01-01")) + wage_growth = float(nawi(f"{year - 1}-01-01")) / float( + nawi(f"{year - 2}-01-01") + ) + updated_value = round_amount(previous_value * wage_growth, rounding) + parameter.update( + period=f"year:{year}-01-01:1", + value=updated_value, + ) + + +def create_wage_indexed_core_thresholds_reform( + *, + start_year: int = 2035, + end_year: int = 2100, +): + from policyengine_us.model_api import Reform + + def modify_parameters(parameters): + nawi = parameters.gov.ssa.nawi + roots = [ + parameters.gov.irs.income.bracket.thresholds, + parameters.gov.irs.deductions.standard.amount, + parameters.gov.irs.deductions.standard.aged_or_blind.amount, + parameters.gov.irs.capital_gains.thresholds, + parameters.gov.irs.income.amt.brackets, + parameters.gov.irs.income.amt.exemption.amount, + parameters.gov.irs.income.amt.exemption.phase_out.start, + parameters.gov.irs.income.amt.exemption.separate_limit, + ] + + seen = set() + for root in roots: + for parameter in iter_updatable_parameters(root): + if parameter.name in seen: + continue + seen.add(parameter.name) + apply_wage_growth_to_parameter( + parameter, + nawi=nawi, + start_year=start_year, + end_year=end_year, + ) + return parameters + + class reform(Reform): + def apply(self): + self.modify_parameters(modify_parameters) + + return reform + + +def create_wage_indexed_full_irs_uprating_reform( + *, + start_year: int = 2035, + end_year: int = 2100, +): + from policyengine_us.model_api import Reform + + def modify_parameters(parameters): + nawi = parameters.gov.ssa.nawi + seen = set() + for parameter in iter_updatable_parameters( + parameters.gov.irs, + uprating_parameter="gov.irs.uprating", + ): + if parameter.name in seen: + continue + seen.add(parameter.name) + apply_wage_growth_to_parameter( + parameter, + nawi=nawi, + start_year=start_year, + end_year=end_year, + ) + return parameters + + class reform(Reform): + def apply(self): + self.modify_parameters(modify_parameters) + + return reform + + +def get_long_run_tax_assumption_metadata( + name: str, + *, + end_year: int, +) -> dict[str, Any]: + if name != TRUSTEES_CORE_THRESHOLD_ASSUMPTION["name"]: + raise ValueError(f"Unknown long-run tax assumption: {name}") + + metadata = dict(TRUSTEES_CORE_THRESHOLD_ASSUMPTION) + metadata["end_year"] = int(end_year) + return metadata + diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index 3236e9be1..2c05f8673 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -33,6 +33,7 @@ aggregate_age_targets, aggregate_household_age_matrix, build_age_bins, + validate_projected_social_security_cap, ) from policyengine_us_data.datasets.cps.long_term.ssa_data import ( available_long_term_target_sources, @@ -51,6 +52,8 @@ ) from policyengine_us_data.datasets.cps.long_term.prototype_synthetic_2100_support import ( SyntheticCandidate, + _compose_role_donor_rows_to_target, + build_role_composite_calibration_blueprint, build_role_donor_composites, summarize_realized_clone_translation, ) @@ -114,8 +117,8 @@ def test_named_profile_lookup(): assert profile.use_greg is False assert profile.use_ss is True assert profile.use_payroll is True - assert profile.use_tob is False - assert profile.benchmark_tob is True + assert profile.use_tob is True + assert profile.benchmark_tob is False assert profile.use_h6_reform is False assert profile.max_negative_weight_pct == 0.0 assert profile.approximate_windows[0].age_bucket_size == 5 @@ -447,6 +450,69 @@ def test_age_bin_helpers_preserve_population_totals(): assert aggregated_target_matrix[:, 1].sum() == pytest.approx((y * 2).sum()) +def test_validate_projected_social_security_cap_rejects_flat_tail(): + from types import SimpleNamespace + + def accessor(year: int): + cap = 254_400.0 if year >= 2035 else 186_000.0 + return SimpleNamespace( + gov=SimpleNamespace( + irs=SimpleNamespace( + payroll=SimpleNamespace( + social_security=SimpleNamespace(cap=cap) + ) + ) + ) + ) + + with pytest.raises(RuntimeError, match="flat after 2035"): + validate_projected_social_security_cap(accessor, 2100) + + +def test_role_composite_calibration_blueprint_reweights_clone_priors(): + report = { + "target_year": 2100, + "clone_household_reports": [ + { + "clone_household_id": 1001, + "target_head_age": 70, + "target_spouse_age": 68, + "target_dependent_ages": [12], + "target_ss_total": 20_000.0, + "target_payroll_total": 50_000.0, + "per_clone_weight_share_pct": 60.0, + }, + { + "clone_household_id": 1002, + "target_head_age": 80, + "target_spouse_age": None, + "target_dependent_ages": [], + "target_ss_total": 30_000.0, + "target_payroll_total": 10_000.0, + "per_clone_weight_share_pct": 40.0, + }, + ], + } + baseline_weights = np.array([10.0, 20.0, 30.0], dtype=float) + blueprint = build_role_composite_calibration_blueprint( + report, + year=2100, + age_bins=build_age_bins(n_ages=86, bucket_size=5), + hh_id_to_idx={999: 0, 1001: 1, 1002: 2}, + baseline_weights=baseline_weights, + base_weight_scale=0.5, + ) + + assert blueprint is not None + assert blueprint["baseline_weights"].tolist() == pytest.approx([5.0, 36.0, 24.0]) + assert blueprint["ss_overrides"] == {1: 20_000.0, 2: 30_000.0} + assert blueprint["payroll_overrides"] == {1: 50_000.0, 2: 10_000.0} + assert blueprint["age_overrides"][1].sum() == pytest.approx(3.0) + assert blueprint["age_overrides"][2].sum() == pytest.approx(1.0) + assert blueprint["summary"]["clone_household_count"] == 2 + assert blueprint["summary"]["base_weight_scale"] == pytest.approx(0.5) + + def test_legacy_flags_map_to_named_profile(): profile = build_profile_from_flags( use_greg=False, @@ -748,7 +814,7 @@ def test_manifest_updates_and_rejects_profile_mismatch(tmp_path): assert rebuilt["years"] == [2026, 2027] -def test_benchmarked_tob_does_not_affect_quality_classification(): +def test_hard_target_tob_affects_quality_classification(): profile = get_profile("ss-payroll-tob") quality = classify_calibration_quality( { @@ -758,8 +824,6 @@ def test_benchmarked_tob_does_not_affect_quality_classification(): "constraints": { "ss_total": {"pct_error": 0.0}, "payroll_total": {"pct_error": 0.0}, - }, - "benchmarks": { "oasdi_tob": {"pct_error": 12.0}, "hi_tob": {"pct_error": -9.0}, }, @@ -767,7 +831,7 @@ def test_benchmarked_tob_does_not_affect_quality_classification(): profile, year=2035, ) - assert quality == "exact" + assert quality == "aggregate" def test_entropy_calibration_produces_nonnegative_weights_and_hits_targets(): @@ -1096,6 +1160,55 @@ def test_manifest_persists_support_augmentation_metadata(tmp_path): ) +def test_manifest_persists_tax_assumption_metadata(tmp_path): + profile = get_profile("ss-payroll-tob") + audit = { + "method_used": "entropy", + "fell_back_to_ipf": False, + "age_max_pct_error": 0.0, + "negative_weight_pct": 0.0, + "positive_weight_count": 70000, + "effective_sample_size": 5000.0, + "top_10_weight_share_pct": 1.5, + "top_100_weight_share_pct": 10.0, + "max_constraint_pct_error": 0.0, + "constraints": {}, + "validation_passed": True, + "validation_issues": [], + } + tax_assumption = { + "name": "trustees-core-thresholds-v1", + "start_year": 2035, + "end_year": 2100, + } + + year_2100 = tmp_path / "2100.h5" + year_2100.write_text("", encoding="utf-8") + metadata_path = write_year_metadata( + year_2100, + year=2100, + base_dataset_path="test.h5", + profile=profile.to_dict(), + calibration_audit=audit, + tax_assumption=tax_assumption, + ) + manifest_path = update_dataset_manifest( + tmp_path, + year=2100, + h5_path=year_2100, + metadata_path=metadata_path, + base_dataset_path="test.h5", + profile=profile.to_dict(), + calibration_audit=audit, + tax_assumption=tax_assumption, + ) + + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + assert metadata["tax_assumption"]["name"] == "trustees-core-thresholds-v1" + assert manifest["tax_assumption"]["end_year"] == 2100 + + def test_write_support_augmentation_report(tmp_path): report = { "name": "donor-backed-composite-v1", @@ -1109,6 +1222,18 @@ def test_write_support_augmentation_report(tmp_path): assert loaded["clone_household_reports"][0]["clone_household_id"] == 1001 +def test_write_support_augmentation_report_custom_filename(tmp_path): + report = {"name": "dynamic-augmentation", "target_year": 2090} + report_path = write_support_augmentation_report( + tmp_path, + report, + filename="support_augmentation_report_2090.json", + ) + assert report_path == tmp_path / "support_augmentation_report_2090.json" + loaded = json.loads(report_path.read_text(encoding="utf-8")) + assert loaded["target_year"] == 2090 + + def test_summarize_realized_clone_translation_matches_toy_clone(): import pandas as pd @@ -1138,3 +1263,47 @@ def test_summarize_realized_clone_translation_matches_toy_clone(): assert summary["aggregate_ss_pct_error"] == pytest.approx(0.0) assert summary["aggregate_payroll_pct_error"] == pytest.approx(0.0) assert summary["per_clone"][0]["realized_ages"] == [70, 68] + + +def test_compose_role_donor_rows_falls_back_for_missing_dependents(): + import pandas as pd + + df = pd.DataFrame(_toy_support_dataframe()) + enriched = df.copy() + enriched["__pe_payroll_uprating_factor"] = 2.0 + enriched["__pe_ss_uprating_factor"] = 3.0 + + older_rows = enriched[enriched["person_tax_unit_id__2024"] == 201].copy() + worker_rows = enriched[enriched["person_tax_unit_id__2024"] == 301].copy() + candidate = SyntheticCandidate( + archetype="older_plus_prime_worker_family_role_donor", + head_age=80, + spouse_age=60, + dependent_ages=(12,), + head_wages=0.0, + spouse_wages=100_000.0, + head_ss=60_000.0, + spouse_ss=0.0, + pension_income=0.0, + dividend_income=0.0, + ) + clone_df, _ = _compose_role_donor_rows_to_target( + older_rows, + worker_rows, + base_year=2024, + target_candidate=candidate, + ss_scale=3.0, + earnings_scale=2.0, + id_counters={ + "household": 100, + "family": 200, + "tax_unit": 300, + "spm_unit": 400, + "marital_unit": 500, + "person": 600, + }, + clone_weight_scale=0.1, + clone_weight_divisor=1, + ) + assert clone_df is not None + assert sorted(clone_df["age__2024"].astype(int).tolist()) == [12, 60, 80] From d287237c1c4b785ef2cd7e11f71becce2037cfbd Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 2 Apr 2026 09:03:38 -0400 Subject: [PATCH 30/31] Add post-OBBBA OACT target source --- .../build_long_term_target_sources.py | 141 ++++++++++++++++++ .../oact_2025_08_05_provisional.csv | 77 ++++++++++ .../oasdi_oact_20250805_nominal_delta.csv | 76 ++++++++++ .../long_term_target_sources/sources.json | 17 +++ .../test_long_term_calibration_contract.py | 12 ++ 5 files changed, 323 insertions(+) create mode 100644 policyengine_us_data/datasets/cps/long_term/build_long_term_target_sources.py create mode 100644 policyengine_us_data/storage/long_term_target_sources/oact_2025_08_05_provisional.csv create mode 100644 policyengine_us_data/storage/long_term_target_sources/oasdi_oact_20250805_nominal_delta.csv diff --git a/policyengine_us_data/datasets/cps/long_term/build_long_term_target_sources.py b/policyengine_us_data/datasets/cps/long_term/build_long_term_target_sources.py new file mode 100644 index 000000000..5fe6db658 --- /dev/null +++ b/policyengine_us_data/datasets/cps/long_term/build_long_term_target_sources.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import pandas as pd +from policyengine_us_data.storage import STORAGE_FOLDER + + +SOURCES_DIR = STORAGE_FOLDER / "long_term_target_sources" +LEGACY_AUX_PATH = STORAGE_FOLDER / "social_security_aux.csv" +OACT_DELTA_PATH = SOURCES_DIR / "oasdi_oact_20250805_nominal_delta.csv" +TRUSTEES_OUTPUT_PATH = SOURCES_DIR / "trustees_2025_current_law.csv" +OACT_OUTPUT_PATH = SOURCES_DIR / "oact_2025_08_05_provisional.csv" +MANIFEST_PATH = SOURCES_DIR / "sources.json" + + +def build_trustees_source() -> pd.DataFrame: + trustees = pd.read_csv(LEGACY_AUX_PATH).copy() + trustees.to_csv(TRUSTEES_OUTPUT_PATH, index=False) + return trustees + + +def build_oact_source(trustees: pd.DataFrame) -> pd.DataFrame: + delta = pd.read_csv(OACT_DELTA_PATH).copy() + if 2100 not in set(delta.year): + delta = pd.concat( + [ + delta, + pd.DataFrame( + { + "year": [2100], + "oasdi_nominal_delta_billions": [ + float(delta.iloc[-1]["oasdi_nominal_delta_billions"]) + ], + } + ), + ], + ignore_index=True, + ) + + merged = trustees.merge(delta, on="year", how="left") + if merged["oasdi_nominal_delta_billions"].isna().any(): + missing_years = merged.loc[ + merged["oasdi_nominal_delta_billions"].isna(), "year" + ].tolist() + raise ValueError( + f"Missing OACT OASDI deltas for years: {missing_years}" + ) + + merged["oasdi_tob_billions_nominal_usd"] = ( + merged["oasdi_tob_billions_nominal_usd"] + + merged["oasdi_nominal_delta_billions"] + ) + ratio = ( + merged["oasdi_tob_billions_nominal_usd"] + / trustees["oasdi_tob_billions_nominal_usd"] + ) + merged["hi_tob_billions_nominal_usd"] = ( + trustees["hi_tob_billions_nominal_usd"] * ratio + ) + merged["oasdi_tob_pct_of_taxable_payroll"] = ( + merged["oasdi_tob_billions_nominal_usd"] + / merged["taxable_payroll_in_billion_nominal_usd"] + * 100 + ) + merged = merged.drop(columns=["oasdi_nominal_delta_billions"]) + merged.to_csv(OACT_OUTPUT_PATH, index=False) + return merged + + +def write_manifest() -> None: + manifest = { + "default_source": "trustees_2025_current_law", + "sources": { + "trustees_2025_current_law": { + "name": "trustees_2025_current_law", + "file": TRUSTEES_OUTPUT_PATH.name, + "type": "trustees_current_law", + "description": ( + "2025 Trustees current-law baseline used by the legacy " + "long-term calibration stack." + ), + "source_urls": [ + "https://www.ssa.gov/oact/tr/2025/lrIndex.html", + "https://www.ssa.gov/oact/solvency/provisions/tables/table_run133.html", + ], + "notes": [ + "Generated from social_security_aux.csv for explicit source selection.", + ], + }, + "oact_2025_08_05_provisional": { + "name": "oact_2025_08_05_provisional", + "file": OACT_OUTPUT_PATH.name, + "type": "oact_override", + "description": ( + "Post-OBBBA SSA OACT baseline overlay with provisional HI " + "bridge for long-term calibration experiments." + ), + "source_urls": [ + "https://www.ssa.gov/OACT/solvency/RWyden_20250805.pdf", + "https://www.ssa.gov/oact/tr/2025/lrIndex.html", + ], + "notes": [ + "OASDI TOB nominal deltas are taken from the August 5, 2025 OACT letter.", + "2100 OASDI delta is carried forward from 2099 because the published delta table ends at 2099.", + "HI TOB series is provisional: it applies the same percentage change as OASDI TOB to preserve the OASDI/HI share split until a published annual HI replacement series is available.", + ], + "derived_from": "trustees_2025_current_law", + "hi_method": "match_oasdi_pct_change", + }, + }, + } + MANIFEST_PATH.write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Build explicit long-term target source packages.", + ) + return parser.parse_args() + + +def main() -> int: + parse_args() + SOURCES_DIR.mkdir(parents=True, exist_ok=True) + trustees = build_trustees_source() + build_oact_source(trustees) + write_manifest() + print(f"Wrote {TRUSTEES_OUTPUT_PATH}") + print(f"Wrote {OACT_OUTPUT_PATH}") + print(f"Wrote {MANIFEST_PATH}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/policyengine_us_data/storage/long_term_target_sources/oact_2025_08_05_provisional.csv b/policyengine_us_data/storage/long_term_target_sources/oact_2025_08_05_provisional.csv new file mode 100644 index 000000000..c73592790 --- /dev/null +++ b/policyengine_us_data/storage/long_term_target_sources/oact_2025_08_05_provisional.csv @@ -0,0 +1,77 @@ +year,oasdi_cost_in_billion_2025_usd,cpi_w_intermediate,oasdi_cost_in_billion_nominal_usd,taxable_payroll_in_billion_nominal_usd,h6_income_rate_change,oasdi_tob_pct_of_taxable_payroll,oasdi_tob_billions_nominal_usd,hi_tob_billions_nominal_usd +2025,1609,100.0,1609.0,10621,0.0,0.5370464174748141,57.0397,38.304600179716786 +2026,1660,102.49,1701.334,11129,0.0,0.5444343606793063,60.59009999999999,41.186853903042184 +2027,1715,104.95,1799.8925,11627,0.0,0.5637885955104498,65.5517,48.095931607707655 +2028,1763,107.47,1894.6961,12159,0.0,0.5852512542149847,71.16069999999999,52.62050146630209 +2029,1810,110.05,1991.905,12696,0.0,0.6373660995589162,80.92,60.16481033396345 +2030,1856,112.69,2091.5264,13239,0.0,0.6689644232948108,88.5642,65.69059107415735 +2031,1903,115.4,2196.062,13798,0.0,0.6983896216843021,96.3638,71.59442052614531 +2032,1947,118.17,2300.7699,14380,0.0,0.7315159944367178,105.19200000000001,78.27395310947746 +2033,1991,121.0,2409.11,14987,0.0,0.7539080536464937,112.9882,85.07501730492007 +2034,2032,123.91,2517.8512,15594,0.0,0.7854726176734642,122.4866,92.24537468746983 +2035,2073,126.88,2630.2224,16205,0.0,0.7981795742054921,129.345,99.04787709554665 +2036,2114,129.93,2746.7202,16825,0.0,0.8113372956909362,136.5075,105.50951161765427 +2037,2155,133.04,2867.012,17465,0.0,0.8149527626681936,142.3315,111.90107428672101 +2038,2194,136.24,2989.1056,18132,0.0,0.8278976395323187,150.11440000000002,118.48385077067687 +2039,2233,139.51,3115.2583,18819,0.0,0.8412598969126945,158.3167,125.23374749266783 +2040,2270,142.86,3242.922,19532,0.0,0.8550112635674789,167.0008,132.13926202500207 +2041,2306,146.28,3373.2168,20269,0.0,0.8581015343628201,173.92860000000002,138.9165967996406 +2042,2342,149.79,3508.0818,21035,0.0,0.861559305918707,181.229,145.92339950639513 +2043,2378,153.39,3647.6142,21828,0.0,0.8748671431189299,190.966,153.25093791653404 +2044,2415,157.07,3793.2405,22653,0.0,0.8780448505716681,198.90349999999998,160.68590468556505 +2045,2452,160.84,3943.7968,23507,-0.07,0.891084357850853,209.46720000000002,168.5513909261071 +2046,2488,164.7,4097.736,24391,-0.12,0.8939920462465664,218.05360000000002,176.52152325652904 +2047,2527,168.65,4261.7855,25313,-0.18,0.896791372022281,227.0048,184.82403098539618 +2048,2567,172.7,4433.209,26270,-0.23,0.9094746859535591,238.91899999999998,193.69560564165155 +2049,2609,176.85,4614.0165,27263,-0.27,0.9120459964053844,248.65109999999999,202.86065543346197 +2050,2652,181.09,4802.5068,28300,-0.32,0.9241696113074204,261.53999999999996,216.92146837816398 +2051,2696,185.44,4999.4624,29376,-0.36,0.9265550108932461,272.1848,227.17710670657155 +2052,2743,189.89,5208.6827,30494,-0.4,0.9385144618613499,286.1906,238.0205394497212 +2053,2792,194.44,5428.7648,31661,-0.43,0.9407280250150026,297.84389999999996,249.44686394215998 +2054,2842,199.11,5658.7062,32869,-1.0,0.9525388664090784,313.09,261.4423901244333 +2055,2895,203.89,5902.6155,34124,-1.01,0.9645774235142422,329.1524,274.24655199383494 +2056,2950,208.78,6159.01,35432,-1.01,0.966254233461278,342.3632,287.5390828182389 +2057,3007,213.79,6428.6653,36790,-1.02,0.9778689861375374,359.758,301.67833439393695 +2058,3066,218.93,6712.3938,38201,-1.03,0.9896869191905971,378.07030000000003,316.6402407536223 +2059,3125,224.18,7005.625,39670,-1.04,1.0011797327955634,397.168,332.06147451086855 +2060,3184,229.56,7309.1904,41196,-1.04,1.0026177298766872,413.0384,347.92088498883385 +2061,3243,235.07,7623.3201,42782,-1.05,1.0140035528960778,433.81100000000004,364.3585166503047 +2062,3303,240.71,7950.6513,44429,-1.06,1.025563033154021,455.6474,381.58588180463465 +2063,3362,246.49,8286.9938,46136,-1.06,1.0268371770417895,473.7416,399.16745314590827 +2064,3422,252.4,8637.128,47902,-1.07,1.038268548286084,497.35139999999996,417.5120418767757 +2065,3483,258.46,9002.1618,49733,-1.07,1.0394367924717993,516.9431,436.4536805413806 +2066,3544,264.66,9379.5504,51631,-1.08,1.050754004377215,542.5147999999999,456.34830162327114 +2067,3607,271.02,9775.6914,53598,-1.09,1.0618273069890667,569.1182,477.000102597283 +2068,3670,277.52,10184.984,55637,-1.09,1.0630395240577315,591.4433,498.52457562710316 +2069,3735,284.18,10614.123,57746,-1.1,1.0741973470023898,620.306,521.0794614282611 +2070,3801,291.0,11060.91,59930,-1.1,1.0751376606040381,644.33,544.3715195000228 +2071,3867,297.99,11523.2733,62196,-1.11,1.0862042575085213,675.5756,568.6915353613309 +2072,3934,305.14,12004.2076,64543,-1.12,1.0970695505322035,708.0816000000001,593.8535429669257 +2073,4002,312.46,12504.6492,66975,-1.12,1.0980515117581187,735.42,620.2432200447928 +2074,4071,319.96,13025.5716,69501,-1.13,1.1089931080128343,770.7613,647.5960347191796 +2075,4139,327.64,13561.0196,72131,-1.13,1.1097590495071465,800.4802999999999,675.8383507115801 +2076,4206,335.5,14111.13,74862,-1.14,1.1206310277577407,838.9268,704.8661033532109 +2077,4273,343.55,14679.8915,77698,-1.14,1.1214667044196762,871.3572,734.6364395845724 +2078,4339,351.8,15264.602,80650,-1.14,1.1222690638561688,905.11,765.1167787494154 +2079,4403,360.24,15861.3672,83727,-1.15,1.1330401184802992,948.6605,796.6188317534055 +2080,4467,368.89,16478.3163,86933,-1.15,1.1337806126557235,985.6295,828.9385545905168 +2081,4530,377.74,17111.622,90268,-1.15,1.1344906279080074,1024.082,862.0944955793473 +2082,4593,386.81,17766.1833,93749,-1.15,1.1351731751805352,1064.2134999999998,896.2517963396247 +2083,4655,396.09,18437.9895,97381,-1.15,1.1358288577853997,1106.0815,931.4300349246773 +2084,4716,405.6,19128.096,101163,-1.15,1.1365563496535294,1149.7745,967.5207711639717 +2085,4775,415.33,19832.0075,105104,-1.15,1.1371555792358043,1195.196,1003.870068209045 +2086,4833,425.3,20554.749,109217,-1.14,1.127730847761795,1231.6737999999998,1041.3446863480701 +2087,4891,435.51,21300.7941,113504,-1.14,1.1283704539047081,1280.7456,1079.6614728458444 +2088,4948,445.96,22066.1008,117973,-1.14,1.1288957642850483,1331.7922,1119.009014012573 +2089,5006,456.66,22860.3996,122629,-1.13,1.1194804654690163,1372.8076999999998,1159.6737769420636 +2090,5064,467.62,23680.2768,127477,-1.13,1.1199589729912063,1427.6901,1201.7744537535523 +2091,5125,478.84,24540.55,132518,-1.13,1.1204918577098961,1484.8534000000002,1245.7033522589752 +2092,5188,490.34,25438.8392,137764,-1.12,1.1109265120060392,1530.4568,1291.5522519728356 +2093,5254,502.1,26380.334,143215,-1.12,1.1113416890688823,1591.608,1339.5725734983864 +2094,5323,514.16,27368.7368,148876,-1.12,1.1117380907600958,1655.1112,1390.0478250557512 +2095,5396,526.49,28409.4004,154754,-1.12,1.1121811390981817,1721.1448,1443.5177747752655 +2096,5472,539.13,29501.1936,160855,-1.12,1.1125398650958938,1789.576,1499.387856294711 +2097,5551,552.07,30645.4057,167185,-1.11,1.102882136555313,1843.8535,1557.8239986469107 +2098,5633,565.32,31844.4756,173750,-1.11,1.1032086330935251,1916.825,1619.4655487069804 +2099,5719,578.89,33106.7191,180557,-1.12,1.113464667667274,2010.4384,1684.301218679657 +2100,5809,592.78,34434.5902,187614,-1.12,1.113710490688328,2089.4768,1751.6199945298017 diff --git a/policyengine_us_data/storage/long_term_target_sources/oasdi_oact_20250805_nominal_delta.csv b/policyengine_us_data/storage/long_term_target_sources/oasdi_oact_20250805_nominal_delta.csv new file mode 100644 index 000000000..c5aabfd19 --- /dev/null +++ b/policyengine_us_data/storage/long_term_target_sources/oasdi_oact_20250805_nominal_delta.csv @@ -0,0 +1,76 @@ +year,oasdi_nominal_delta_billions +2025,-3.5 +2026,-16.2 +2027,-17.0 +2028,-17.6 +2029,-14.3 +2030,-14.7 +2031,-15.4 +2032,-15.6 +2033,-15.9 +2034,-16.3 +2035,-16.5 +2036,-16.6 +2037,-16.6 +2038,-16.7 +2039,-16.7 +2040,-16.6 +2041,-16.6 +2042,-16.5 +2043,-16.4 +2044,-16.3 +2045,-16.2 +2046,-16.1 +2047,-16.0 +2048,-15.9 +2049,-15.8 +2050,-15.8 +2051,-15.7 +2052,-15.7 +2053,-15.6 +2054,-15.6 +2055,-15.5 +2056,-15.5 +2057,-15.5 +2058,-15.4 +2059,-15.4 +2060,-15.4 +2061,-15.4 +2062,-15.3 +2063,-15.3 +2064,-15.2 +2065,-15.2 +2066,-15.1 +2067,-15.1 +2068,-15.0 +2069,-14.9 +2070,-14.9 +2071,-14.8 +2072,-14.8 +2073,-14.7 +2074,-14.6 +2075,-14.6 +2076,-14.5 +2077,-14.4 +2078,-14.3 +2079,-14.2 +2080,-14.1 +2081,-14.0 +2082,-13.9 +2083,-13.8 +2084,-13.6 +2085,-13.5 +2086,-13.4 +2087,-13.2 +2088,-13.1 +2089,-12.9 +2090,-12.8 +2091,-12.6 +2092,-12.5 +2093,-12.4 +2094,-12.3 +2095,-12.1 +2096,-12.0 +2097,-11.9 +2098,-11.8 +2099,-11.8 diff --git a/policyengine_us_data/storage/long_term_target_sources/sources.json b/policyengine_us_data/storage/long_term_target_sources/sources.json index 3b1932697..62dd15e1a 100644 --- a/policyengine_us_data/storage/long_term_target_sources/sources.json +++ b/policyengine_us_data/storage/long_term_target_sources/sources.json @@ -1,6 +1,23 @@ { "default_source": "trustees_2025_current_law", "sources": { + "oact_2025_08_05_provisional": { + "derived_from": "trustees_2025_current_law", + "description": "Post-OBBBA SSA OACT baseline overlay with provisional HI bridge for long-term calibration experiments.", + "file": "oact_2025_08_05_provisional.csv", + "hi_method": "match_oasdi_pct_change", + "name": "oact_2025_08_05_provisional", + "notes": [ + "OASDI TOB nominal deltas are taken from the August 5, 2025 OACT letter.", + "2100 OASDI delta is carried forward from 2099 because the published delta table ends at 2099.", + "HI TOB series is provisional: it applies the same percentage change as OASDI TOB to preserve the OASDI/HI share split until a published annual HI replacement series is available." + ], + "source_urls": [ + "https://www.ssa.gov/OACT/solvency/RWyden_20250805.pdf", + "https://www.ssa.gov/oact/tr/2025/lrIndex.html" + ], + "type": "oact_override" + }, "trustees_2025_current_law": { "description": "2025 Trustees current-law baseline used by the legacy long-term calibration stack.", "file": "trustees_2025_current_law.csv", diff --git a/policyengine_us_data/tests/test_long_term_calibration_contract.py b/policyengine_us_data/tests/test_long_term_calibration_contract.py index 2c05f8673..1c4b33047 100644 --- a/policyengine_us_data/tests/test_long_term_calibration_contract.py +++ b/policyengine_us_data/tests/test_long_term_calibration_contract.py @@ -38,6 +38,7 @@ from policyengine_us_data.datasets.cps.long_term.ssa_data import ( available_long_term_target_sources, describe_long_term_target_source, + load_oasdi_tob_projections, load_taxable_payroll_projections, ) from policyengine_us_data.datasets.cps.long_term.support_augmentation import ( @@ -977,6 +978,7 @@ def test_nonnegative_feasibility_diagnostic_distinguishes_feasible_and_infeasibl def test_long_term_target_sources_are_available_and_distinct(): sources = available_long_term_target_sources() assert "trustees_2025_current_law" in sources + assert "oact_2025_08_05_provisional" in sources trustees = describe_long_term_target_source("trustees_2025_current_law") assert trustees["file"] == "trustees_2025_current_law.csv" @@ -987,6 +989,16 @@ def test_long_term_target_sources_are_available_and_distinct(): ) assert payroll_2026 == pytest.approx(11_129_000_000_000.0) + trustees_oasdi_2026 = load_oasdi_tob_projections( + 2026, + source_name="trustees_2025_current_law", + ) + oact_oasdi_2026 = load_oasdi_tob_projections( + 2026, + source_name="oact_2025_08_05_provisional", + ) + assert oact_oasdi_2026 < trustees_oasdi_2026 + def test_normalize_metadata_backfills_validation_passed(): metadata = normalize_metadata( From f6182503a132f0d920a3c1aee776ae83fdff335a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 2 Apr 2026 21:16:39 -0400 Subject: [PATCH 31/31] Format long-run calibration files --- .../long_term/assess_calibration_frontier.py | 13 +- .../long_term/assess_publishable_horizon.py | 12 +- .../benchmark_trustees_bracket_indexing.py | 6 +- .../build_long_term_target_sources.py | 4 +- .../datasets/cps/long_term/calibration.py | 71 ++--- .../cps/long_term/calibration_artifacts.py | 19 +- .../cps/long_term/calibration_profiles.py | 7 +- .../cps/long_term/compare_tob_shares.py | 4 +- ...agnose_support_augmentation_translation.py | 4 +- .../evaluate_support_augmentation.py | 8 +- .../profile_support_concentration.py | 30 ++- .../cps/long_term/projection_utils.py | 11 +- .../prototype_synthetic_2100_support.py | 250 ++++++++++++------ .../cps/long_term/run_household_projection.py | 85 ++---- .../long_term/summarize_calibration_runs.py | 21 +- .../cps/long_term/support_augmentation.py | 116 ++++---- .../datasets/cps/long_term/tax_assumptions.py | 1 - .../test_long_term_calibration_contract.py | 61 +++-- 18 files changed, 382 insertions(+), 341 deletions(-) diff --git a/policyengine_us_data/datasets/cps/long_term/assess_calibration_frontier.py b/policyengine_us_data/datasets/cps/long_term/assess_calibration_frontier.py index 04f024850..f9fefd8fe 100644 --- a/policyengine_us_data/datasets/cps/long_term/assess_calibration_frontier.py +++ b/policyengine_us_data/datasets/cps/long_term/assess_calibration_frontier.py @@ -86,13 +86,17 @@ def reorder_to_households(values, order, n_households: int) -> np.ndarray: return ordered -def build_constraint_inputs(year: int, hh_id_to_idx: dict, n_households: int, profile) -> dict: +def build_constraint_inputs( + year: int, hh_id_to_idx: dict, n_households: int, profile +) -> dict: sim = Microsimulation(dataset=BASE_DATASET_PATH) if profile.use_h6_reform: raise NotImplementedError( "Frontier assessment for H6-enabled profiles is not yet implemented." ) - household_ids = sim.calculate("household_id", period=year, map_to="household").values + household_ids = sim.calculate( + "household_id", period=year, map_to="household" + ).values if len(household_ids) != n_households: raise ValueError( f"Household count mismatch for {year}: {len(household_ids)} vs {n_households}" @@ -226,10 +230,7 @@ def main() -> int: } rows.append(row) best_case_display = "n/a" if best_case is None else f"{best_case:.3f}%" - print( - f"{year}: best-case max error {best_case_display} -> " - f"{row['quality']}" - ) + print(f"{year}: best-case max error {best_case_display} -> {row['quality']}") if args.output: output_path = Path(args.output) diff --git a/policyengine_us_data/datasets/cps/long_term/assess_publishable_horizon.py b/policyengine_us_data/datasets/cps/long_term/assess_publishable_horizon.py index 3c482243c..dcf5cd08c 100644 --- a/policyengine_us_data/datasets/cps/long_term/assess_publishable_horizon.py +++ b/policyengine_us_data/datasets/cps/long_term/assess_publishable_horizon.py @@ -119,9 +119,7 @@ def benchmark_tob_values( else (oasdi_achieved - oasdi_target) / oasdi_target * 100 ), "hi_tob_benchmark_pct_error": ( - 0.0 - if hi_target == 0 - else (hi_achieved - hi_target) / hi_target * 100 + 0.0 if hi_target == 0 else (hi_achieved - hi_target) / hi_target * 100 ), } @@ -287,9 +285,7 @@ def assess_years( "validation_issues": str(error), "runtime_error": str(error), } - best_case_match = re.search( - r"([0-9.]+)%\s*>\s*([0-9.]+)%", str(error) - ) + best_case_match = re.search(r"([0-9.]+)%\s*>\s*([0-9.]+)%", str(error)) if best_case_match: row["reported_best_case_constraint_error_pct"] = float( best_case_match.group(1) @@ -315,7 +311,9 @@ def assess_years( h6_income_values=None, h6_revenue_target=None, oasdi_tob_values=oasdi_tob_values if profile.use_tob else None, - oasdi_tob_target=load_oasdi_tob_projections(year) if profile.use_tob else None, + oasdi_tob_target=load_oasdi_tob_projections(year) + if profile.use_tob + else None, hi_tob_values=hi_tob_values if profile.use_tob else None, hi_tob_target=load_hi_tob_projections(year) if profile.use_tob else None, ) diff --git a/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py b/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py index 5d44edbe6..0929c1c2e 100644 --- a/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py +++ b/policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py @@ -45,9 +45,7 @@ def modify_parameters(parameters): bracket_node = thresholds.get_child(bracket) for filing_status in FILING_STATUSES: parameter = bracket_node.get_child(filing_status) - interval = float( - parameter.metadata["uprating"]["rounding"]["interval"] - ) + interval = float(parameter.metadata["uprating"]["rounding"]["interval"]) for year in range(start_year, end_year + 1): previous_value = float(parameter(f"{year - 1}-01-01")) @@ -66,6 +64,8 @@ def apply(self): self.modify_parameters(modify_parameters) return reform + + def _coerce_h5_path(raw: str) -> Path: path = Path(raw).expanduser() if path.is_dir(): diff --git a/policyengine_us_data/datasets/cps/long_term/build_long_term_target_sources.py b/policyengine_us_data/datasets/cps/long_term/build_long_term_target_sources.py index 5fe6db658..20655c354 100644 --- a/policyengine_us_data/datasets/cps/long_term/build_long_term_target_sources.py +++ b/policyengine_us_data/datasets/cps/long_term/build_long_term_target_sources.py @@ -45,9 +45,7 @@ def build_oact_source(trustees: pd.DataFrame) -> pd.DataFrame: missing_years = merged.loc[ merged["oasdi_nominal_delta_billions"].isna(), "year" ].tolist() - raise ValueError( - f"Missing OACT OASDI deltas for years: {missing_years}" - ) + raise ValueError(f"Missing OACT OASDI deltas for years: {missing_years}") merged["oasdi_tob_billions_nominal_usd"] = ( merged["oasdi_tob_billions_nominal_usd"] diff --git a/policyengine_us_data/datasets/cps/long_term/calibration.py b/policyengine_us_data/datasets/cps/long_term/calibration.py index 238350b4d..92dc0a4ec 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration.py @@ -279,9 +279,7 @@ def calibrate_entropy( baseline_weights = np.asarray(baseline_weights, dtype=float) gram = A_scaled.T @ (baseline_weights[:, None] * A_scaled) gram += np.eye(gram.shape[0]) * 1e-12 - beta0 = np.linalg.solve( - gram, targets_scaled - (A_scaled.T @ baseline_weights) - ) + beta0 = np.linalg.solve(gram, targets_scaled - (A_scaled.T @ baseline_weights)) def objective_gradient_hessian(beta): eta = np.clip(A_scaled @ beta, -700, 700) @@ -313,11 +311,7 @@ def _cached_ogh(z): return None _, gradient, _ = objective_gradient_hessian(result.x) max_error = float( - np.max( - 100 - * np.abs(gradient) - / np.maximum(np.abs(targets_scaled), 1e-12) - ) + np.max(100 * np.abs(gradient) / np.maximum(np.abs(targets_scaled), 1e-12)) ) if max_error > tol * 100: return None @@ -343,11 +337,7 @@ def infeasibility_error(prefix): for iterations in range(1, max_iters + 1): objective, gradient, hessian = objective_gradient_hessian(beta) final_max_error = float( - np.max( - 100 - * np.abs(gradient) - / np.maximum(np.abs(targets_scaled), 1e-12) - ) + np.max(100 * np.abs(gradient) / np.maximum(np.abs(targets_scaled), 1e-12)) ) if final_max_error <= tol * 100: break @@ -361,8 +351,8 @@ def infeasibility_error(prefix): step = 1.0 while step >= 1e-8: candidate = beta - step * delta - candidate_objective, candidate_gradient, _ = ( - objective_gradient_hessian(candidate) + candidate_objective, candidate_gradient, _ = objective_gradient_hessian( + candidate ) candidate_max_error = float( np.max( @@ -477,9 +467,7 @@ def objective_and_gradient(z): exp_eta = np.exp(eta) weights = baseline_weights * exp_eta achieved = A_scaled.T @ weights - objective = float( - np.sum(weights) + upper_bounds @ alpha - lower_bounds @ gamma - ) + objective = float(np.sum(weights) + upper_bounds @ alpha - lower_bounds @ gamma) gradient = np.concatenate( [ upper_bounds - achieved, @@ -561,9 +549,7 @@ def objective_with_gradient(z): ) if best_result is None or best_weights is None: - raise RuntimeError( - "Approximate bounded entropy calibration did not run." - ) + raise RuntimeError("Approximate bounded entropy calibration did not run.") raise RuntimeError( "Approximate bounded entropy calibration failed: " @@ -610,8 +596,7 @@ def densify_lp_solution( lam = (lo + hi) / 2.0 candidate_weights = (1.0 - lam) * lp_weights + lam * baseline_weights candidate_error_pct = float( - np.max(np.abs(A_scaled.T @ candidate_weights - targets_scaled)) - * 100 + np.max(np.abs(A_scaled.T @ candidate_weights - targets_scaled)) * 100 ) if candidate_error_pct <= max_constraint_error_pct + 1e-6: best_lambda = lam @@ -670,14 +655,11 @@ def calibrate_lp_minimax( A = aux_df.to_numpy(dtype=float) targets = np.array(list(controls.values()), dtype=float) - feasibility = assess_nonnegative_feasibility( - A, targets, return_weights=True - ) + feasibility = assess_nonnegative_feasibility(A, targets, return_weights=True) weights = feasibility.get("weights") if not feasibility["success"] or weights is None: raise RuntimeError( - "Approximate nonnegative calibration failed: " - f"{feasibility['message']}" + f"Approximate nonnegative calibration failed: {feasibility['message']}" ) return np.asarray(weights, dtype=float), 1, feasibility @@ -706,9 +688,7 @@ def assess_nonnegative_feasibility(A, targets, *, return_weights=False): b_rel = targets / scales constraint_matrix = sparse.csr_matrix(A_rel) - epsilon_column = sparse.csc_matrix( - np.ones((constraint_matrix.shape[0], 1)) - ) + epsilon_column = sparse.csc_matrix(np.ones((constraint_matrix.shape[0], 1))) A_ub = sparse.vstack( [ sparse.hstack([constraint_matrix, -epsilon_column]), @@ -886,9 +866,7 @@ def calibrate_weights( hi_tob_target=hi_tob_target, n_ages=n_ages, ) - approximate_error_pct = float( - feasibility["best_case_max_pct_error"] - ) + approximate_error_pct = float(feasibility["best_case_max_pct_error"]) if approximate_error_pct <= max(tol * 100, 1e-6): audit["lp_fallback_used"] = True audit["approximation_method"] = "lp_minimax_exact" @@ -1019,9 +997,7 @@ def build_calibration_audit( ): achieved_ages = X.T @ weights age_errors = ( - np.abs(achieved_ages - y_target) - / np.maximum(np.abs(y_target), 1e-10) - * 100 + np.abs(achieved_ages - y_target) / np.maximum(np.abs(y_target), 1e-10) * 100 ) neg_mask = weights < 0 @@ -1031,12 +1007,8 @@ def build_calibration_audit( abs_weight_sum = float(np.sum(np.abs(weights))) if weight_sum > 0: sorted_weights = np.sort(weights) - top_10_weight_share_pct = float( - sorted_weights[-10:].sum() / weight_sum * 100 - ) - top_100_weight_share_pct = float( - sorted_weights[-100:].sum() / weight_sum * 100 - ) + top_10_weight_share_pct = float(sorted_weights[-10:].sum() / weight_sum * 100) + top_100_weight_share_pct = float(sorted_weights[-100:].sum() / weight_sum * 100) else: top_10_weight_share_pct = 0.0 top_100_weight_share_pct = 0.0 @@ -1051,9 +1023,7 @@ def build_calibration_audit( { "age_max_pct_error": float(age_errors.max()), "negative_weight_count": int(neg_mask.sum()), - "negative_weight_household_pct": float( - 100 * neg_mask.sum() / len(weights) - ), + "negative_weight_household_pct": float(100 * neg_mask.sum() / len(weights)), "negative_weight_pct": ( float(100 * negative_values.sum() / abs_weight_sum) if abs_weight_sum > 0 @@ -1063,9 +1033,7 @@ def build_calibration_audit( float(negative_values.max()) if negative_values.size else 0.0 ), "positive_weight_count": int(positive_mask.sum()), - "positive_weight_pct": float( - 100 * positive_mask.sum() / len(weights) - ), + "positive_weight_pct": float(100 * positive_mask.sum() / len(weights)), "effective_sample_size": effective_sample_size, "top_10_weight_share_pct": top_10_weight_share_pct, "top_100_weight_share_pct": top_100_weight_share_pct, @@ -1097,10 +1065,7 @@ def build_calibration_audit( if audit["constraints"]: audit["max_constraint_pct_error"] = float( - max( - abs(stats["pct_error"]) - for stats in audit["constraints"].values() - ) + max(abs(stats["pct_error"]) for stats in audit["constraints"].values()) ) return audit diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py index 610a7a83e..3dcb86f2e 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py @@ -39,10 +39,7 @@ def normalize_metadata(metadata: dict[str, Any]) -> dict[str, Any]: if "max_constraint_pct_error" not in audit: audit["max_constraint_pct_error"] = float( max( - ( - abs(stats.get("pct_error", 0.0)) - for stats in constraints.values() - ), + (abs(stats.get("pct_error", 0.0)) for stats in constraints.values()), default=0.0, ) ) @@ -186,11 +183,9 @@ def update_dataset_manifest( ) manifest_profile = json.loads(json.dumps(manifest["profile"])) if manifest_profile != profile: - if manifest_profile.get("name") == profile.get( - "name" - ) and manifest_profile.get("calibration_method") == profile.get( + if manifest_profile.get("name") == profile.get("name") and manifest_profile.get( "calibration_method" - ): + ) == profile.get("calibration_method"): manifest["profile"] = profile else: raise ValueError( @@ -230,17 +225,13 @@ def update_dataset_manifest( "method_used": calibration_audit.get("method_used"), "fell_back_to_ipf": calibration_audit.get("fell_back_to_ipf"), "age_max_pct_error": calibration_audit.get("age_max_pct_error"), - "max_constraint_pct_error": calibration_audit.get( - "max_constraint_pct_error" - ), + "max_constraint_pct_error": calibration_audit.get("max_constraint_pct_error"), "negative_weight_pct": calibration_audit.get("negative_weight_pct"), "negative_weight_household_pct": calibration_audit.get( "negative_weight_household_pct" ), "validation_passed": calibration_audit.get("validation_passed"), - "validation_issue_count": len( - calibration_audit.get("validation_issues", []) - ), + "validation_issue_count": len(calibration_audit.get("validation_issues", [])), } year_set = {int(value) for value in manifest.get("years", [])} diff --git a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py index b093bc882..15c91f8a6 100644 --- a/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py +++ b/policyengine_us_data/datasets/cps/long_term/calibration_profiles.py @@ -203,7 +203,9 @@ def get_profile(name: str) -> CalibrationProfile: return NAMED_PROFILES[name] except KeyError as error: valid = ", ".join(sorted(NAMED_PROFILES)) - raise ValueError(f"Unknown calibration profile '{name}'. Valid profiles: {valid}") from error + raise ValueError( + f"Unknown calibration profile '{name}'. Valid profiles: {valid}" + ) from error def approximate_window_for_year( @@ -459,8 +461,7 @@ def _collect_threshold_issues( and ess < min_effective_sample_size ): issues.append( - f"Effective sample size {ess:.3f} is below " - f"{min_effective_sample_size:.3f}" + f"Effective sample size {ess:.3f} is below {min_effective_sample_size:.3f}" ) top_10_share = audit.get("top_10_weight_share_pct") diff --git a/policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py b/policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py index 9ab4104e8..91b0e13a5 100644 --- a/policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py +++ b/policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py @@ -83,9 +83,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "paths", nargs="+", - help=( - "Metadata files or directories containing *.metadata.json sidecars." - ), + help=("Metadata files or directories containing *.metadata.json sidecars."), ) parser.add_argument( "--format", diff --git a/policyengine_us_data/datasets/cps/long_term/diagnose_support_augmentation_translation.py b/policyengine_us_data/datasets/cps/long_term/diagnose_support_augmentation_translation.py index 505a647d4..3230380f3 100644 --- a/policyengine_us_data/datasets/cps/long_term/diagnose_support_augmentation_translation.py +++ b/policyengine_us_data/datasets/cps/long_term/diagnose_support_augmentation_translation.py @@ -36,7 +36,9 @@ def main() -> int: "report_file named in metadata, or /support_augmentation_report.json." ), ) - parser.add_argument("--year", type=int, required=True, help="Output year to inspect.") + parser.add_argument( + "--year", type=int, required=True, help="Output year to inspect." + ) parser.add_argument( "--age-bucket-size", type=int, diff --git a/policyengine_us_data/datasets/cps/long_term/evaluate_support_augmentation.py b/policyengine_us_data/datasets/cps/long_term/evaluate_support_augmentation.py index b70624378..a90c23382 100644 --- a/policyengine_us_data/datasets/cps/long_term/evaluate_support_augmentation.py +++ b/policyengine_us_data/datasets/cps/long_term/evaluate_support_augmentation.py @@ -49,9 +49,7 @@ def _evaluate_dataset( approximate_window = approximate_window_for_year(profile, year) age_bucket_size = ( - approximate_window.age_bucket_size - if approximate_window is not None - else None + approximate_window.age_bucket_size if approximate_window is not None else None ) if age_bucket_size and age_bucket_size > 1: age_bins = build_age_bins(n_ages=n_ages, bucket_size=age_bucket_size) @@ -68,7 +66,9 @@ def _evaluate_dataset( ss_values = None ss_target = None if profile.use_ss: - ss_values = sim.calculate("social_security", period=year, map_to="household").values + ss_values = sim.calculate( + "social_security", period=year, map_to="household" + ).values ss_target = load_ssa_benefit_projections(year) payroll_values = None diff --git a/policyengine_us_data/datasets/cps/long_term/profile_support_concentration.py b/policyengine_us_data/datasets/cps/long_term/profile_support_concentration.py index b6102dad5..b79ded0ff 100644 --- a/policyengine_us_data/datasets/cps/long_term/profile_support_concentration.py +++ b/policyengine_us_data/datasets/cps/long_term/profile_support_concentration.py @@ -30,7 +30,11 @@ def _build_household_lookup(household_ids: np.ndarray) -> tuple[np.ndarray, np.n return sorted_ids, order -def _household_index(sorted_household_ids: np.ndarray, order: np.ndarray, person_household_ids: np.ndarray) -> np.ndarray: +def _household_index( + sorted_household_ids: np.ndarray, + order: np.ndarray, + person_household_ids: np.ndarray, +) -> np.ndarray: positions = np.searchsorted(sorted_household_ids, person_household_ids) if np.any(positions >= len(sorted_household_ids)): raise ValueError("Person household ids exceed household id support") @@ -43,8 +47,12 @@ def _household_index(sorted_household_ids: np.ndarray, order: np.ndarray, person def _load_year(path: Path, year: int) -> dict[str, np.ndarray]: with h5py.File(path, "r") as store: household_ids = _read_year_array(store, "household_id", year).astype(np.int64) - household_weights = _read_year_array(store, "household_weight", year).astype(float) - person_household_ids = _read_year_array(store, "person_household_id", year).astype(np.int64) + household_weights = _read_year_array(store, "household_weight", year).astype( + float + ) + person_household_ids = _read_year_array( + store, "person_household_id", year + ).astype(np.int64) ages = _read_year_array(store, "age", year).astype(float) payroll = np.zeros_like(ages, dtype=float) for component in PAYROLL_COMPONENTS: @@ -117,8 +125,12 @@ def profile_support(path: Path, year: int, *, top_n: int) -> dict[str, object]: "positive_household_count": int(positive_mask.sum()), "positive_household_pct": float(positive_mask.mean() * 100), "effective_sample_size": _effective_sample_size(household_weights), - "top_10_weight_share_pct": float(sorted_weights[-10:].sum() / household_weights.sum() * 100), - "top_100_weight_share_pct": float(sorted_weights[-100:].sum() / household_weights.sum() * 100), + "top_10_weight_share_pct": float( + sorted_weights[-10:].sum() / household_weights.sum() * 100 + ), + "top_100_weight_share_pct": float( + sorted_weights[-100:].sum() / household_weights.sum() * 100 + ), "weighted_nonworking_share_pct": float( person_weights[overall_nonworking].sum() / person_weights.sum() * 100 ), @@ -137,9 +149,13 @@ def main() -> None: parser = argparse.ArgumentParser( description="Profile late-year support concentration in projected household datasets." ) - parser.add_argument("dataset", type=Path, help="Projected year-specific H5 dataset.") + parser.add_argument( + "dataset", type=Path, help="Projected year-specific H5 dataset." + ) parser.add_argument("year", type=int, help="Projection year stored in the dataset.") - parser.add_argument("--top-n", type=int, default=20, help="Number of top households to emit.") + parser.add_argument( + "--top-n", type=int, default=20, help="Number of top households to emit." + ) args = parser.parse_args() report = profile_support(args.dataset, args.year, top_n=args.top_n) diff --git a/policyengine_us_data/datasets/cps/long_term/projection_utils.py b/policyengine_us_data/datasets/cps/long_term/projection_utils.py index 7c28005f6..1c0baf5bc 100644 --- a/policyengine_us_data/datasets/cps/long_term/projection_utils.py +++ b/policyengine_us_data/datasets/cps/long_term/projection_utils.py @@ -22,9 +22,7 @@ def validate_projected_social_security_cap( through 2100. If the payroll cap flattens after the reference year, the late-year taxable payroll problem becomes mechanically distorted. """ - current_cap = float( - parameter_accessor(year).gov.irs.payroll.social_security.cap - ) + current_cap = float(parameter_accessor(year).gov.irs.payroll.social_security.cap) reference_cap = float( parameter_accessor(reference_year).gov.irs.payroll.social_security.cap ) @@ -111,12 +109,7 @@ def aggregate_age_targets(targets, age_bins): dtype=float, ) - return np.vstack( - [ - targets[start:end, :].sum(axis=0) - for start, end in age_bins - ] - ) + return np.vstack([targets[start:end, :].sum(axis=0) for start, end in age_bins]) def get_pseudo_input_variables(sim): diff --git a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py index 84339fbc3..12ab5b793 100644 --- a/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py +++ b/policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py @@ -412,7 +412,12 @@ def classify_archetype( if older_head: if dependent_count > 0 and positive_payroll: return "older_plus_prime_worker_family" - if has_spouse and spouse_age is not None and spouse_age < 65 and positive_payroll: + if ( + has_spouse + and spouse_age is not None + and spouse_age < 65 + and positive_payroll + ): return "mixed_retiree_worker_couple" if has_spouse and older_spouse: if positive_payroll and positive_ss: @@ -451,9 +456,7 @@ def build_tax_unit_summary( "is_dependent": sim.calculate( "is_tax_unit_dependent", period=period ).values, - "social_security": sim.calculate( - "social_security", period=period - ).values, + "social_security": sim.calculate("social_security", period=period).values, "payroll": ( sim.calculate( "taxable_earnings_for_social_security", period=period @@ -469,9 +472,9 @@ def build_tax_unit_summary( "taxable_pension_income", period=period ).values, "person_weight": input_df[f"person_weight__{period}"].astype(float).values, - "household_weight": input_df[ - f"household_weight__{period}" - ].astype(float).values, + "household_weight": input_df[f"household_weight__{period}"] + .astype(float) + .values, } ) @@ -484,11 +487,11 @@ def build_tax_unit_summary( head_age = float( heads["age"].iloc[0] if not heads.empty - else adults["age"].max() if not adults.empty else group["age"].max() - ) - spouse_age = ( - float(spouses["age"].iloc[0]) if not spouses.empty else None + else adults["age"].max() + if not adults.empty + else group["age"].max() ) + spouse_age = float(spouses["age"].iloc[0]) if not spouses.empty else None dependent_count = int(len(dependents)) adult_count = int((group["age"] >= 18).sum()) dependent_ages = tuple(sorted(int(age) for age in dependents["age"].tolist())) @@ -585,7 +588,9 @@ def attach_person_uprating_factors( def load_base_aggregates(base_dataset: str) -> dict[str, float]: sim = Microsimulation(dataset=base_dataset) - household_series = sim.calculate("household_id", period=BASE_YEAR, map_to="household") + household_series = sim.calculate( + "household_id", period=BASE_YEAR, map_to="household" + ) weights = household_series.weights.values.astype(float) ss = sim.calculate("social_security", period=BASE_YEAR, map_to="household").values payroll = ( @@ -659,7 +664,9 @@ def allocate_taxable_payroll_wages( return 0.0, min(total_taxable_payroll, payroll_cap) return min(total_taxable_payroll, payroll_cap), 0.0 - total_taxable_payroll = min(total_taxable_payroll, positive_earner_count * payroll_cap) + total_taxable_payroll = min( + total_taxable_payroll, positive_earner_count * payroll_cap + ) preferred_head = total_taxable_payroll * float(payroll_split[0]) lower = max(0.0, total_taxable_payroll - payroll_cap) upper = min(payroll_cap, total_taxable_payroll) @@ -775,31 +782,50 @@ def generate_synthetic_candidates( for payroll_total in payroll_levels: for payroll_scale in template.payroll_scale_factors: for pension_income in pension_levels: - for pension_scale in template.pension_scale_factors: + for ( + pension_scale + ) in template.pension_scale_factors: for dividend_income in dividend_levels: - for dividend_scale in template.dividend_scale_factors: - scaled_ss_total = ss_total * ss_scale - scaled_payroll_total = payroll_total * payroll_scale - scaled_pension_income = pension_income * pension_scale - scaled_dividend_income = dividend_income * dividend_scale + for ( + dividend_scale + ) in template.dividend_scale_factors: + scaled_ss_total = ( + ss_total * ss_scale + ) + scaled_payroll_total = ( + payroll_total * payroll_scale + ) + scaled_pension_income = ( + pension_income * pension_scale + ) + scaled_dividend_income = ( + dividend_income * dividend_scale + ) head_ss = ( - scaled_ss_total * template.ss_split[0] + scaled_ss_total + * template.ss_split[0] ) spouse_ss = ( - scaled_ss_total * template.ss_split[1] + scaled_ss_total + * template.ss_split[1] ) - head_wages, spouse_wages = allocate_taxable_payroll_wages( - scaled_payroll_total, - template.payroll_split, - payroll_cap, - has_spouse=spouse_age is not None, + head_wages, spouse_wages = ( + allocate_taxable_payroll_wages( + scaled_payroll_total, + template.payroll_split, + payroll_cap, + has_spouse=spouse_age + is not None, + ) ) candidates.append( SyntheticCandidate( archetype=template.name, head_age=head_age, spouse_age=spouse_age, - dependent_ages=tuple(dependent_ages), + dependent_ages=tuple( + dependent_ages + ), head_wages=head_wages, spouse_wages=spouse_wages, head_ss=head_ss, @@ -983,7 +1009,9 @@ def _clone_report_record( int(older_donor_row["tax_unit_id"]) if older_donor_row is not None else None ), "worker_donor_tax_unit_id": ( - int(worker_donor_row["tax_unit_id"]) if worker_donor_row is not None else None + int(worker_donor_row["tax_unit_id"]) + if worker_donor_row is not None + else None ), "older_donor_distance": ( float(older_donor_row["distance"]) @@ -1059,7 +1087,9 @@ def summarize_realized_clone_translation( "realized_ages": realized_ages, "realized_ss_total": realized_ss_total, "realized_payroll_total": realized_payroll_total, - "realized_household_weight": float(realized_row["household_weight_proxy"]), + "realized_household_weight": float( + realized_row["household_weight_proxy"] + ), "age_bucket_l1": float(np.abs(realized_vector - target_vector).sum()), "ss_pct_error": ( 0.0 @@ -1145,7 +1175,9 @@ def build_synthetic_constraint_problem( age_targets = load_ssa_age_projections(start_year=year, end_year=year) age_bins = build_age_bins(n_ages=age_targets.shape[0], bucket_size=5) aggregated_age_targets = aggregate_age_targets(age_targets, age_bins)[:, 0] - X = np.vstack([age_bucket_vector(candidate.ages(), age_bins) for candidate in candidates]) + X = np.vstack( + [age_bucket_vector(candidate.ages(), age_bins) for candidate in candidates] + ) ss_values = np.array([candidate.ss_total for candidate in candidates], dtype=float) payroll_values = np.array( [candidate.taxable_payroll_total(payroll_cap) for candidate in candidates], @@ -1170,7 +1202,9 @@ def build_synthetic_constraint_problem( } -def build_constraint_matrix(problem: dict[str, object]) -> tuple[np.ndarray, np.ndarray]: +def build_constraint_matrix( + problem: dict[str, object], +) -> tuple[np.ndarray, np.ndarray]: constraint_matrix = np.column_stack( [ problem["X"], @@ -1222,7 +1256,10 @@ def _target_head_ss_share(candidate: SyntheticCandidate) -> float: def _target_worker_age(candidate: SyntheticCandidate) -> float: - if candidate.spouse_age is not None and candidate.spouse_wages > candidate.head_wages: + if ( + candidate.spouse_age is not None + and candidate.spouse_wages > candidate.head_wages + ): return float(candidate.spouse_age) return float(candidate.head_age) @@ -1256,8 +1293,7 @@ def match_older_role_donors( + np.abs( np.log1p(subset["scaled_ss_total"]) - np.log1p(target_candidate.ss_total) ) - + 0.5 - * np.abs(donor_head_ss_share - target_head_ss_share) + + 0.5 * np.abs(donor_head_ss_share - target_head_ss_share) + 0.15 * subset["archetype"] .isin( @@ -1282,9 +1318,7 @@ def match_worker_role_donors( *, donors_per_target: int, ) -> pd.DataFrame: - subset = scaled_actual_summary[ - scaled_actual_summary["payroll_total"] > 0 - ].copy() + subset = scaled_actual_summary[scaled_actual_summary["payroll_total"] > 0].copy() target_dependent_count = len(target_candidate.dependent_ages) target_spouse_present = target_candidate.spouse_age is not None if target_dependent_count > 0: @@ -1309,8 +1343,7 @@ def match_worker_role_donors( np.log1p(subset["scaled_payroll_total"]) - np.log1p(target_candidate.payroll_total) ) - + 0.5 - * np.abs(donor_head_payroll_share - target_head_payroll_share) + + 0.5 * np.abs(donor_head_payroll_share - target_head_payroll_share) + 0.25 * subset["spouse_present"].ne(target_spouse_present).astype(float) ) return subset.nsmallest(donors_per_target, "distance").copy() @@ -1555,7 +1588,9 @@ def match_real_donors_for_target( subset = scaled_actual_summary[ scaled_actual_summary["spouse_present"].eq(target_spouse_present) & scaled_actual_summary["adult_count"].eq(target_adult_count) - & scaled_actual_summary["dependent_count"].eq(int(target_row["dependent_count"])) + & scaled_actual_summary["dependent_count"].eq( + int(target_row["dependent_count"]) + ) ].copy() if subset.empty: subset = scaled_actual_summary[ @@ -1565,10 +1600,7 @@ def match_real_donors_for_target( subset = scaled_actual_summary.copy() subset["distance"] = ( (subset["head_age"] - target_row["head_age"]).abs() / 5.0 - + ( - subset["spouse_age_filled"] - target_spouse_age - ).abs() - / 5.0 + + (subset["spouse_age_filled"] - target_spouse_age).abs() / 5.0 + (subset["dependent_count"] - target_row["dependent_count"]).abs() * 0.75 + np.abs( np.log1p(subset["scaled_payroll_total"]) @@ -1629,10 +1661,14 @@ def summarize_donor_probe( ) ) donor_df = pd.concat(donor_matches, ignore_index=True) - nearest_only = donor_df.sort_values( - ["target_candidate_idx", "distance"], - ascending=[True, True], - ).groupby("target_candidate_idx", as_index=False).first() + nearest_only = ( + donor_df.sort_values( + ["target_candidate_idx", "distance"], + ascending=[True, True], + ) + .groupby("target_candidate_idx", as_index=False) + .first() + ) distance_summary = { "median_best_distance": float(nearest_only["distance"].median()), "targets_with_best_distance_le_1": int((nearest_only["distance"] <= 1.0).sum()), @@ -1756,10 +1792,14 @@ def _clone_tax_unit_rows_to_target( person_id_col = _period_column(PERSON_ID_COLUMN, base_year) adults = donor_rows[donor_rows[age_col] >= 18].sort_values(age_col, ascending=False) - dependents = donor_rows[donor_rows[age_col] < 18].sort_values(age_col, ascending=False) + dependents = donor_rows[donor_rows[age_col] < 18].sort_values( + age_col, ascending=False + ) target_has_spouse = target_candidate.spouse_age is not None target_adult_count = 2 if target_has_spouse else 1 - if len(adults) < target_adult_count or len(dependents) != len(target_candidate.dependent_ages): + if len(adults) < target_adult_count or len(dependents) != len( + target_candidate.dependent_ages + ): return None, id_counters cloned = donor_rows.copy() @@ -1771,7 +1811,9 @@ def _clone_tax_unit_rows_to_target( for raw_column in columns: column = _period_column(raw_column, base_year) if column in cloned.columns: - cloned[column] = entity_id if entity_name != "household" else household_id + cloned[column] = ( + entity_id if entity_name != "household" else household_id + ) cloned[_period_column("household_id", base_year)] = household_id cloned[_period_column("person_household_id", base_year)] = household_id @@ -1806,8 +1848,12 @@ def _clone_tax_unit_rows_to_target( for dep_idx, dep_age in zip(dependent_indices, target_candidate.dependent_ages): cloned.loc[dep_idx, age_col] = float(dep_age) - payroll_columns = tuple(_period_column(component, base_year) for component in PAYROLL_COMPONENTS) - ss_columns = tuple(_period_column(component, base_year) for component in SS_COMPONENTS) + payroll_columns = tuple( + _period_column(component, base_year) for component in PAYROLL_COMPONENTS + ) + ss_columns = tuple( + _period_column(component, base_year) for component in SS_COMPONENTS + ) qbi_col = _period_column("w2_wages_from_qualified_business", base_year) target_head_payroll = _target_base_total_for_row( @@ -1942,7 +1988,9 @@ def _dependent_rows(df: pd.DataFrame | None) -> pd.DataFrame: selected_rows: list[pd.Series] = [] head_target_older = target_candidate.head_ss > 0 or target_candidate.head_age >= 65 - head_source_rows = older_adults if head_target_older and not older_adults.empty else worker_adults + head_source_rows = ( + older_adults if head_target_older and not older_adults.empty else worker_adults + ) if head_source_rows.empty: return None, id_counters head_row = head_source_rows.iloc[0].copy() @@ -1953,15 +2001,25 @@ def _dependent_rows(df: pd.DataFrame | None) -> pd.DataFrame: if target_candidate.spouse_age >= 65 and len(older_adults) >= 2: spouse_row = older_adults.iloc[1].copy() elif not worker_adults.empty: - worker_candidates = worker_adults.iloc[1:] if worker_adults.index[0] == head_row.name else worker_adults + worker_candidates = ( + worker_adults.iloc[1:] + if worker_adults.index[0] == head_row.name + else worker_adults + ) if worker_candidates.empty: worker_candidates = worker_adults - spouse_idx = (worker_candidates[age_col] - target_candidate.spouse_age).abs().idxmin() + spouse_idx = ( + (worker_candidates[age_col] - target_candidate.spouse_age) + .abs() + .idxmin() + ) spouse_row = worker_candidates.loc[spouse_idx].copy() elif len(older_adults) >= 2: spouse_row = older_adults.iloc[1].copy() if spouse_row is None: - fallback_spouse_pool = worker_adults if not worker_adults.empty else older_adults + fallback_spouse_pool = ( + worker_adults if not worker_adults.empty else older_adults + ) if fallback_spouse_pool.empty: return None, id_counters spouse_row = fallback_spouse_pool.iloc[0].copy() @@ -1972,13 +2030,15 @@ def _dependent_rows(df: pd.DataFrame | None) -> pd.DataFrame: if not dependent_rows: fallback_source = None if worker_donor_rows is not None and not worker_donor_rows.empty: - fallback_source = worker_donor_rows.sort_values(age_col, ascending=True).iloc[ - 0 - ].copy() + fallback_source = ( + worker_donor_rows.sort_values(age_col, ascending=True) + .iloc[0] + .copy() + ) elif older_donor_rows is not None and not older_donor_rows.empty: - fallback_source = older_donor_rows.sort_values(age_col, ascending=True).iloc[ - 0 - ].copy() + fallback_source = ( + older_donor_rows.sort_values(age_col, ascending=True).iloc[0].copy() + ) if fallback_source is None: return None, id_counters dependent_rows = [fallback_source.copy()] @@ -1997,7 +2057,9 @@ def _dependent_rows(df: pd.DataFrame | None) -> pd.DataFrame: for raw_column in columns: column = _period_column(raw_column, base_year) if column in cloned.columns: - cloned[column] = entity_id if entity_name != "household" else household_id + cloned[column] = ( + entity_id if entity_name != "household" else household_id + ) cloned[_period_column("household_id", base_year)] = household_id cloned[_period_column("person_household_id", base_year)] = household_id @@ -2035,8 +2097,12 @@ def _dependent_rows(df: pd.DataFrame | None) -> pd.DataFrame: for dep_idx, dep_age in zip(dependent_indices, target_candidate.dependent_ages): cloned.loc[dep_idx, age_col] = float(dep_age) - payroll_columns = tuple(_period_column(component, base_year) for component in PAYROLL_COMPONENTS) - ss_columns = tuple(_period_column(component, base_year) for component in SS_COMPONENTS) + payroll_columns = tuple( + _period_column(component, base_year) for component in PAYROLL_COMPONENTS + ) + ss_columns = tuple( + _period_column(component, base_year) for component in SS_COMPONENTS + ) qbi_col = _period_column("w2_wages_from_qualified_business", base_year) target_head_payroll = float(target_candidate.head_wages) / max( @@ -2146,7 +2212,9 @@ def build_donor_backed_augmented_input_dataframe( entity_name: _next_entity_id(input_df[_period_column(columns[0], base_year)]) for entity_name, columns in ENTITY_ID_COLUMNS.items() } - id_counters["person"] = _next_entity_id(input_df[_period_column(PERSON_ID_COLUMN, base_year)]) + id_counters["person"] = _next_entity_id( + input_df[_period_column(PERSON_ID_COLUMN, base_year)] + ) clone_frames = [] clone_household_reports = [] @@ -2160,7 +2228,9 @@ def build_donor_backed_augmented_input_dataframe( scaled_actual, donors_per_target=donors_per_target, ) - usable = donor_matches[donor_matches["distance"] <= max_distance_for_clone].copy() + usable = donor_matches[ + donor_matches["distance"] <= max_distance_for_clone + ].copy() if usable.empty: skipped_targets.append( { @@ -2174,7 +2244,9 @@ def build_donor_backed_augmented_input_dataframe( successful_clone_count = 0 for _, donor_row in usable.iterrows(): - donor_rows = input_df[input_df[tax_unit_id_col] == int(donor_row["tax_unit_id"])].copy() + donor_rows = input_df[ + input_df[tax_unit_id_col] == int(donor_row["tax_unit_id"]) + ].copy() clone_df, id_counters = _clone_tax_unit_rows_to_target( donor_rows, base_year=base_year, @@ -2220,8 +2292,12 @@ def build_donor_backed_augmented_input_dataframe( "donors_per_target": int(donors_per_target), "max_distance_for_clone": float(max_distance_for_clone), "clone_weight_scale": float(clone_weight_scale), - "base_household_count": int(input_df[_period_column("household_id", base_year)].nunique()), - "augmented_household_count": int(augmented_df[_period_column("household_id", base_year)].nunique()), + "base_household_count": int( + input_df[_period_column("household_id", base_year)].nunique() + ), + "augmented_household_count": int( + augmented_df[_period_column("household_id", base_year)].nunique() + ), "base_person_count": int(len(input_df)), "augmented_person_count": int(len(augmented_df)), "target_reports": target_reports, @@ -2540,7 +2616,9 @@ def build_donor_backed_clones( scaled_actual, donors_per_target=donors_per_target, ) - usable = donor_matches[donor_matches["distance"] <= max_distance_for_clone].copy() + usable = donor_matches[ + donor_matches["distance"] <= max_distance_for_clone + ].copy() if usable.empty: donor_backed_candidates.append(target_candidate) donor_backed_weights.append(float(target_row["synthetic_weight"])) @@ -2586,11 +2664,13 @@ def build_donor_backed_clones( spouse_age=target_candidate.spouse_age, dependent_ages=target_candidate.dependent_ages, head_wages=target_payroll_total * donor_head_payroll_share, - spouse_wages=target_payroll_total * (1.0 - donor_head_payroll_share), + spouse_wages=target_payroll_total + * (1.0 - donor_head_payroll_share), head_ss=target_ss_total * donor_head_ss_share, spouse_ss=target_ss_total * (1.0 - donor_head_ss_share), pension_income=float(donor_row["pension_income"]) * earnings_scale, - dividend_income=float(donor_row["dividend_income"]) * earnings_scale, + dividend_income=float(donor_row["dividend_income"]) + * earnings_scale, ) ) donor_backed_weights.append(per_clone_weight) @@ -2702,9 +2782,7 @@ def solve_synthetic_support( else "lp_minimax" ), "iterations": 1, - "best_case_max_pct_error": float( - dense_info["best_case_max_pct_error"] - ), + "best_case_max_pct_error": float(dense_info["best_case_max_pct_error"]), "requested_max_constraint_error_pct": float(max_constraint_error_pct), "age_bucket_size": 5, "entropy_error": str(error), @@ -2760,7 +2838,9 @@ def summarize_solution_diff( base_total = max(float(base_weights.sum()), 1.0) alt_total = max(float(alt_weights.sum()), 1.0) rows: list[dict[str, object]] = [] - for candidate, base_weight, alt_weight in zip(candidates, base_weights, alt_weights): + for candidate, base_weight, alt_weight in zip( + candidates, base_weights, alt_weights + ): base_share_pct = float(base_weight / base_total * 100) alt_share_pct = float(alt_weight / alt_total * 100) rows.append( @@ -2946,8 +3026,14 @@ def _weighted_mean(group: pd.DataFrame, column: str) -> float: else 0.0, "synthetic_avg_taxable_benefits_proxy_share_pct": float( ( - (candidate_df["taxable_benefits_proxy"] * candidate_df["synthetic_weight"]).sum() - / max((candidate_df["ss_total"] * candidate_df["synthetic_weight"]).sum(), 1.0) + ( + candidate_df["taxable_benefits_proxy"] + * candidate_df["synthetic_weight"] + ).sum() + / max( + (candidate_df["ss_total"] * candidate_df["synthetic_weight"]).sum(), + 1.0, + ) ) * 100 ), @@ -2955,7 +3041,9 @@ def _weighted_mean(group: pd.DataFrame, column: str) -> float: return { "synthetic_candidate_count": int(len(candidate_df)), - "positive_weight_candidate_count": int((candidate_df["synthetic_weight"] > 0).sum()), + "positive_weight_candidate_count": int( + (candidate_df["synthetic_weight"] > 0).sum() + ), "effective_sample_size": effective_sample_size, "top_10_weight_share_pct": top_10_weight_share_pct, "top_20_weight_share_pct": top_20_weight_share_pct, diff --git a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py index f90b51f68..5fd8008de 100644 --- a/policyengine_us_data/datasets/cps/long_term/run_household_projection.py +++ b/policyengine_us_data/datasets/cps/long_term/run_household_projection.py @@ -308,9 +308,7 @@ def create_h6_reform(): if "--support-augmentation-profile" in sys.argv: augmentation_index = sys.argv.index("--support-augmentation-profile") if augmentation_index + 1 >= len(sys.argv): - raise ValueError( - "--support-augmentation-profile requires a profile name" - ) + raise ValueError("--support-augmentation-profile requires a profile name") SUPPORT_AUGMENTATION_PROFILE = sys.argv[augmentation_index + 1] del sys.argv[augmentation_index : augmentation_index + 2] @@ -318,9 +316,7 @@ def create_h6_reform(): if "--support-augmentation-target-year" in sys.argv: target_year_index = sys.argv.index("--support-augmentation-target-year") if target_year_index + 1 >= len(sys.argv): - raise ValueError( - "--support-augmentation-target-year requires a year" - ) + raise ValueError("--support-augmentation-target-year requires a year") SUPPORT_AUGMENTATION_TARGET_YEAR = int(sys.argv[target_year_index + 1]) del sys.argv[target_year_index : target_year_index + 2] @@ -334,9 +330,7 @@ def create_h6_reform(): if "--support-augmentation-start-year" in sys.argv: start_year_index = sys.argv.index("--support-augmentation-start-year") if start_year_index + 1 >= len(sys.argv): - raise ValueError( - "--support-augmentation-start-year requires a year" - ) + raise ValueError("--support-augmentation-start-year requires a year") SUPPORT_AUGMENTATION_START_YEAR = int(sys.argv[start_year_index + 1]) del sys.argv[start_year_index : start_year_index + 2] @@ -344,9 +338,7 @@ def create_h6_reform(): if "--support-augmentation-top-n-targets" in sys.argv: top_n_index = sys.argv.index("--support-augmentation-top-n-targets") if top_n_index + 1 >= len(sys.argv): - raise ValueError( - "--support-augmentation-top-n-targets requires an integer" - ) + raise ValueError("--support-augmentation-top-n-targets requires an integer") SUPPORT_AUGMENTATION_TOP_N_TARGETS = int(sys.argv[top_n_index + 1]) del sys.argv[top_n_index : top_n_index + 2] @@ -354,9 +346,7 @@ def create_h6_reform(): if "--support-augmentation-donors-per-target" in sys.argv: donor_index = sys.argv.index("--support-augmentation-donors-per-target") if donor_index + 1 >= len(sys.argv): - raise ValueError( - "--support-augmentation-donors-per-target requires an integer" - ) + raise ValueError("--support-augmentation-donors-per-target requires an integer") SUPPORT_AUGMENTATION_DONORS_PER_TARGET = int(sys.argv[donor_index + 1]) del sys.argv[donor_index : donor_index + 2] @@ -364,24 +354,16 @@ def create_h6_reform(): if "--support-augmentation-max-distance" in sys.argv: distance_index = sys.argv.index("--support-augmentation-max-distance") if distance_index + 1 >= len(sys.argv): - raise ValueError( - "--support-augmentation-max-distance requires a float" - ) + raise ValueError("--support-augmentation-max-distance requires a float") SUPPORT_AUGMENTATION_MAX_DISTANCE = float(sys.argv[distance_index + 1]) del sys.argv[distance_index : distance_index + 2] SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE = 0.1 if "--support-augmentation-clone-weight-scale" in sys.argv: - weight_scale_index = sys.argv.index( - "--support-augmentation-clone-weight-scale" - ) + weight_scale_index = sys.argv.index("--support-augmentation-clone-weight-scale") if weight_scale_index + 1 >= len(sys.argv): - raise ValueError( - "--support-augmentation-clone-weight-scale requires a float" - ) - SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE = float( - sys.argv[weight_scale_index + 1] - ) + raise ValueError("--support-augmentation-clone-weight-scale requires a float") + SUPPORT_AUGMENTATION_CLONE_WEIGHT_SCALE = float(sys.argv[weight_scale_index + 1]) del sys.argv[weight_scale_index : weight_scale_index + 2] SUPPORT_AUGMENTATION_BLUEPRINT_BASE_WEIGHT_SCALE = 0.5 @@ -464,8 +446,7 @@ def create_h6_reform(): if SUPPORT_AUGMENTATION_PROFILE is not None: if SUPPORT_AUGMENTATION_PROFILE not in SUPPORTED_AUGMENTATION_PROFILES: raise ValueError( - "Unsupported support augmentation profile: " - f"{SUPPORT_AUGMENTATION_PROFILE}" + f"Unsupported support augmentation profile: {SUPPORT_AUGMENTATION_PROFILE}" ) if START_YEAR < SUPPORT_AUGMENTATION_START_YEAR: raise ValueError( @@ -574,10 +555,7 @@ def _compose_reforms(*reforms): if SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR: print(" Support augmentation target year: each run year") else: - print( - " Support augmentation target year: " - f"{SUPPORT_AUGMENTATION_TARGET_YEAR}" - ) + print(f" Support augmentation target year: {SUPPORT_AUGMENTATION_TARGET_YEAR}") print( " Support augmentation blueprint base-weight scale: " f"{SUPPORT_AUGMENTATION_BLUEPRINT_BASE_WEIGHT_SCALE}" @@ -662,9 +640,7 @@ def _build_support_augmentation( "augmented_household_count" ], "base_person_count": augmentation_report["base_person_count"], - "augmented_person_count": augmentation_report[ - "augmented_person_count" - ], + "augmented_person_count": augmentation_report["augmented_person_count"], "clone_household_count": augmentation_report.get( "clone_household_count", 0 ), @@ -692,9 +668,7 @@ def _build_support_augmentation( SUPPORT_AUGMENTATION_BLUEPRINT_BASE_WEIGHT_SCALE ), "report_file": ( - None - if SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR - else report_path.name + None if SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR else report_path.name ), } return augmented_dataset, augmentation_report, year_metadata, manifest_metadata @@ -715,10 +689,8 @@ def _print_support_augmentation_summary(augmentation_report: dict) -> None: " Successful target clones: " f"{sum(report['successful_clone_count'] > 0 for report in augmentation_report['target_reports'])}" ) - print( - " Skipped synthetic targets: " - f"{len(augmentation_report['skipped_targets'])}" - ) + print(f" Skipped synthetic targets: {len(augmentation_report['skipped_targets'])}") + # ========================================================================= # STEP 1: LOAD SSA DEMOGRAPHIC PROJECTIONS @@ -808,7 +780,9 @@ def _print_support_augmentation_summary(augmentation_report: dict) -> None: print("=" * 70) if SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR: - print("\nDynamic augmentation enabled; base support will be used before the activation year and rebuilt per-year after that.") + print( + "\nDynamic augmentation enabled; base support will be used before the activation year and rebuilt per-year after that." + ) sim = Microsimulation(dataset=BASE_DATASET) X, household_ids_unique, hh_id_to_idx = build_household_age_matrix(sim, n_ages) @@ -870,9 +844,7 @@ def _print_support_augmentation_summary(augmentation_report: dict) -> None: _, ) = cached if year in display_years: - print( - f" [DEBUG {year}] Rebuilt support augmentation for run year {year}" - ) + print(f" [DEBUG {year}] Rebuilt support augmentation for run year {year}") _print_support_augmentation_summary(current_support_augmentation_report) sim = Microsimulation( dataset=current_dataset, @@ -933,9 +905,7 @@ def _print_support_augmentation_summary(augmentation_report: dict) -> None: payroll_target = load_taxable_payroll_projections(year) if year in display_years: payroll_baseline = np.sum(payroll_values * baseline_weights) - print( - f" [DEBUG {year}] Payroll cap: ${payroll_cap:,.0f}" - ) + print(f" [DEBUG {year}] Payroll cap: ${payroll_cap:,.0f}") print( f" [DEBUG {year}] Payroll baseline: ${payroll_baseline / 1e9:.1f}B, target: ${payroll_target / 1e9:.1f}B" ) @@ -1015,12 +985,13 @@ def _print_support_augmentation_summary(augmentation_report: dict) -> None: approximate_window = approximate_window_for_year(PROFILE, year) age_bucket_size = ( - approximate_window.age_bucket_size - if approximate_window is not None - else None + approximate_window.age_bucket_size if approximate_window is not None else None ) if age_bucket_size and age_bucket_size > 1: - if SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR and current_support_augmentation_report is not None: + if ( + SUPPORT_AUGMENTATION_ALIGN_TO_RUN_YEAR + and current_support_augmentation_report is not None + ): age_bins = build_age_bins(n_ages=n_ages, bucket_size=age_bucket_size) X_current = aggregate_household_age_matrix(current_X, age_bins) aggregated_target_matrix = aggregate_age_targets(target_matrix, age_bins) @@ -1053,7 +1024,7 @@ def _print_support_augmentation_summary(augmentation_report: dict) -> None: ) blueprint_summary = None if ( - SUPPORT_AUGMENTATION_PROFILE == "donor-backed-composite-v1" + SUPPORT_AUGMENTATION_PROFILE == "donor-backed-composite-v1" and current_support_augmentation_report is not None and year >= SUPPORT_AUGMENTATION_START_YEAR ): @@ -1170,7 +1141,9 @@ def _print_support_augmentation_summary(augmentation_report: dict) -> None: if validation_issues: issue_text = "; ".join(validation_issues) if not ALLOW_VALIDATION_FAILURES: - raise RuntimeError(f"Calibration validation failed for {year}: {issue_text}") + raise RuntimeError( + f"Calibration validation failed for {year}: {issue_text}" + ) print( f" [WARN {year}] Validation issues recorded but not fatal: {issue_text}", file=sys.stderr, diff --git a/policyengine_us_data/datasets/cps/long_term/summarize_calibration_runs.py b/policyengine_us_data/datasets/cps/long_term/summarize_calibration_runs.py index da9535a4e..029d27985 100644 --- a/policyengine_us_data/datasets/cps/long_term/summarize_calibration_runs.py +++ b/policyengine_us_data/datasets/cps/long_term/summarize_calibration_runs.py @@ -29,7 +29,9 @@ def parse_args() -> argparse.Namespace: ) ) parser.add_argument("left", type=Path, help="First projected dataset directory.") - parser.add_argument("right", type=Path, nargs="?", help="Optional second directory to compare.") + parser.add_argument( + "right", type=Path, nargs="?", help="Optional second directory to compare." + ) parser.add_argument( "--years", help="Optional comma-separated list of years to include. Defaults to all years found.", @@ -60,12 +62,11 @@ def metadata_for(directory: Path, year: int) -> dict[str, Any] | None: return normalize_metadata(json.loads(metadata_path.read_text(encoding="utf-8"))) -def support_metrics(directory: Path, year: int, metadata: dict[str, Any] | None, *, profile: bool) -> dict[str, Any]: +def support_metrics( + directory: Path, year: int, metadata: dict[str, Any] | None, *, profile: bool +) -> dict[str, Any]: audit = (metadata or {}).get("calibration_audit", {}) - metrics = { - field: audit.get(field) - for field in SUPPORT_FIELDS - } + metrics = {field: audit.get(field) for field in SUPPORT_FIELDS} if not profile or all(value is not None for value in metrics.values()): return metrics @@ -80,11 +81,15 @@ def support_metrics(directory: Path, year: int, metadata: dict[str, Any] | None, "top_10_weight_share_pct": profiled["top_10_weight_share_pct"], "top_100_weight_share_pct": profiled["top_100_weight_share_pct"], "weighted_nonworking_share_pct": profiled["weighted_nonworking_share_pct"], - "weighted_nonworking_share_85_plus_pct": profiled["weighted_nonworking_share_85_plus_pct"], + "weighted_nonworking_share_85_plus_pct": profiled[ + "weighted_nonworking_share_85_plus_pct" + ], } -def summarize_directory(directory: Path, years: list[int] | None, *, profile: bool) -> dict[int, dict[str, Any]]: +def summarize_directory( + directory: Path, years: list[int] | None, *, profile: bool +) -> dict[int, dict[str, Any]]: if years is None: years = sorted( int(path.name.split(".")[0]) diff --git a/policyengine_us_data/datasets/cps/long_term/support_augmentation.py b/policyengine_us_data/datasets/cps/long_term/support_augmentation.py index 8976a2424..d23bfe3ff 100644 --- a/policyengine_us_data/datasets/cps/long_term/support_augmentation.py +++ b/policyengine_us_data/datasets/cps/long_term/support_augmentation.py @@ -19,9 +19,7 @@ "employment_income_before_lsr", "self_employment_income_before_lsr", ) -PAYROLL_TRANSFER_COMPONENTS = PAYROLL_COMPONENTS + ( - "w2_wages_from_qualified_business", -) +PAYROLL_TRANSFER_COMPONENTS = PAYROLL_COMPONENTS + ("w2_wages_from_qualified_business",) ENTITY_ID_COLUMNS = { "household": ("household_id", "person_household_id"), "family": ("family_id", "person_family_id"), @@ -103,7 +101,7 @@ class SupportAugmentationProfile: | CompositePayrollRule | SinglePersonSyntheticGridRule | MixedAgeAppendRule, - ... + ..., ] @@ -462,7 +460,10 @@ def household_support_summary( {_period_column(component, base_year): "sum" for component in SS_COMPONENTS} ) aggregations.update( - {_period_column(component, base_year): "sum" for component in PAYROLL_COMPONENTS} + { + _period_column(component, base_year): "sum" + for component in PAYROLL_COMPONENTS + } ) summary = ( @@ -481,11 +482,9 @@ def household_support_summary( summary["payroll_total"] = summary[ [_period_column(component, base_year) for component in PAYROLL_COMPONENTS] ].sum(axis=1) - summary["household_size"] = ( - input_df.groupby(household_id_col, sort=False) - [_period_column(PERSON_ID_COLUMN, base_year)] - .count() - ) + summary["household_size"] = input_df.groupby(household_id_col, sort=False)[ + _period_column(PERSON_ID_COLUMN, base_year) + ].count() return summary @@ -673,20 +672,28 @@ def _quantile_pair_households( ) -> list[tuple[int, int]]: if recipient_ids.empty or donor_ids.empty: return [] - recipient_order = summary.loc[recipient_ids].sort_values( - ["ss_total", "baseline_weight", "max_age"] - ).index.to_list() - donor_order = summary.loc[donor_ids].sort_values( - ["payroll_total", "baseline_weight", "max_age"] - ).index.to_list() + recipient_order = ( + summary.loc[recipient_ids] + .sort_values(["ss_total", "baseline_weight", "max_age"]) + .index.to_list() + ) + donor_order = ( + summary.loc[donor_ids] + .sort_values(["payroll_total", "baseline_weight", "max_age"]) + .index.to_list() + ) if len(donor_order) == 1: donor_positions = np.zeros(len(recipient_order), dtype=int) else: - donor_positions = np.linspace( - 0, - len(donor_order) - 1, - num=len(recipient_order), - ).round().astype(int) + donor_positions = ( + np.linspace( + 0, + len(donor_order) - 1, + num=len(recipient_order), + ) + .round() + .astype(int) + ) return [ (int(recipient_household_id), int(donor_order[position])) for recipient_household_id, position in zip(recipient_order, donor_positions) @@ -704,10 +711,9 @@ def _select_payroll_target_row( adults = household_rows[household_rows[age_col] >= 18] if adults.empty: adults = household_rows - existing_payroll = ( - adults.get(employment_col, 0).astype(float) - + adults.get(self_employment_col, 0).astype(float) - ) + existing_payroll = adults.get(employment_col, 0).astype(float) + adults.get( + self_employment_col, 0 + ).astype(float) if existing_payroll.gt(0).any(): return existing_payroll.idxmax() return adults[age_col].astype(float).idxmax() @@ -808,9 +814,11 @@ def synthesize_composite_households( ) original_recipients = pd.unique( - input_df[input_df[household_id_col].isin([recipient for recipient, _ in recipient_pairs])][ - household_id_col - ] + input_df[ + input_df[household_id_col].isin( + [recipient for recipient, _ in recipient_pairs] + ) + ][household_id_col] ) cloned_household_ids = pd.unique(clone_df[household_id_col]) cloned_mapping = { @@ -936,7 +944,9 @@ def synthesize_single_person_grid_households( for target_age in rule.target_ages: for ss_target in ss_targets: for payroll_target in payroll_targets: - base_row = template_records[template_index % len(template_records)].copy() + base_row = template_records[ + template_index % len(template_records) + ].copy() template_index += 1 household_id = next_ids["household"] next_ids["household"] += 1 @@ -951,11 +961,13 @@ def synthesize_single_person_grid_households( base_row[column] = entity_id base_row[household_id_col] = household_id - base_row[_period_column("person_household_id", base_year)] = household_id + base_row[_period_column("person_household_id", base_year)] = ( + household_id + ) base_row[person_id_col] = person_id base_row[age_col] = float(target_age) - base_row[_period_column("social_security_retirement", base_year)] = float( - ss_target + base_row[_period_column("social_security_retirement", base_year)] = ( + float(ss_target) ) base_row[_period_column("social_security_disability", base_year)] = 0.0 base_row[_period_column("social_security_survivors", base_year)] = 0.0 @@ -982,9 +994,7 @@ def synthesize_single_person_grid_households( "rule": rule.name, "template_household_count": int(len(template_ids)), "payroll_donor_household_count": int(len(donor_ids)), - "synthetic_household_count": int( - synthetic_df[household_id_col].nunique() - ), + "synthetic_household_count": int(synthetic_df[household_id_col].nunique()), "synthetic_person_count": int(len(synthetic_df)), "target_age_count": int(len(rule.target_ages)), "ss_grid_size": int(len(rule.ss_quantiles)), @@ -1066,7 +1076,9 @@ def synthesize_mixed_age_households( donor_rows, base_year=base_year, ) - household_weight = float(cloned_household_rows.iloc[0][_period_column("household_weight", base_year)]) + household_weight = float( + cloned_household_rows.iloc[0][_period_column("household_weight", base_year)] + ) person_weight = ( float(cloned_household_rows.iloc[0][person_weight_col]) if person_weight_col in cloned_household_rows.columns @@ -1092,9 +1104,7 @@ def synthesize_mixed_age_households( "rule": rule.name, "recipient_household_count": int(len(recipient_ids)), "donor_household_count": int(len(donor_ids)), - "synthetic_household_count": int( - synthetic_df[household_id_col].nunique() - ), + "synthetic_household_count": int(synthetic_df[household_id_col].nunique()), "synthetic_person_count": int(len(synthetic_df)), "donor_age_shift": rule.donor_age_shift, }, @@ -1152,12 +1162,14 @@ def augment_input_dataframe( continue if isinstance(rule, CompositePayrollRule): - composite_df, id_counters, composite_report = synthesize_composite_households( - input_df, - base_year=base_year, - summary=summary, - rule=rule, - id_counters=id_counters, + composite_df, id_counters, composite_report = ( + synthesize_composite_households( + input_df, + base_year=base_year, + summary=summary, + rule=rule, + id_counters=id_counters, + ) ) clone_frames.append(composite_df) rule_reports.append(composite_report) @@ -1175,12 +1187,14 @@ def augment_input_dataframe( rule_reports.append(mixed_report) continue - synthetic_df, id_counters, synthetic_report = synthesize_single_person_grid_households( - input_df, - base_year=base_year, - summary=summary, - rule=rule, - id_counters=id_counters, + synthetic_df, id_counters, synthetic_report = ( + synthesize_single_person_grid_households( + input_df, + base_year=base_year, + summary=summary, + rule=rule, + id_counters=id_counters, + ) ) clone_frames.append(synthetic_df) rule_reports.append(synthetic_report) diff --git a/policyengine_us_data/datasets/cps/long_term/tax_assumptions.py b/policyengine_us_data/datasets/cps/long_term/tax_assumptions.py index e93f4580e..4e29c825f 100644 --- a/policyengine_us_data/datasets/cps/long_term/tax_assumptions.py +++ b/policyengine_us_data/datasets/cps/long_term/tax_assumptions.py @@ -175,4 +175,3 @@ def get_long_run_tax_assumption_metadata( metadata = dict(TRUSTEES_CORE_THRESHOLD_ASSUMPTION) metadata["end_year"] = int(end_year) return metadata - diff --git a/tests/unit/test_long_term_calibration_contract.py b/tests/unit/test_long_term_calibration_contract.py index 1c4b33047..73c24b4d9 100644 --- a/tests/unit/test_long_term_calibration_contract.py +++ b/tests/unit/test_long_term_calibration_contract.py @@ -226,12 +226,12 @@ def test_support_augmentation_synthesizes_composite_payroll_household(): augmented_df["household_id__2024"].isin(cloned_household_ids) ] assert cloned_rows["age__2024"].max() == pytest.approx(80.0) - assert cloned_rows[ - "social_security_retirement__2024" - ].sum() == pytest.approx(30_000.0) - assert cloned_rows[ - "employment_income_before_lsr__2024" - ].sum() == pytest.approx(37_000.0) + assert cloned_rows["social_security_retirement__2024"].sum() == pytest.approx( + 30_000.0 + ) + assert cloned_rows["employment_income_before_lsr__2024"].sum() == pytest.approx( + 37_000.0 + ) def test_support_augmentation_appends_single_person_synthetic_grid_households(): @@ -288,20 +288,22 @@ def test_support_augmentation_appends_single_person_synthetic_grid_households(): ) assert report["base_household_count"] == 3 assert report["augmented_household_count"] == 5 - synthetic_household_ids = set( - augmented_df["household_id__2024"].unique() - ) - {1, 2, 3} + synthetic_household_ids = set(augmented_df["household_id__2024"].unique()) - { + 1, + 2, + 3, + } assert len(synthetic_household_ids) == 2 synthetic_rows = augmented_df[ augmented_df["household_id__2024"].isin(synthetic_household_ids) ] assert set(synthetic_rows["age__2024"].tolist()) == {77.0, 85.0} - assert set( - synthetic_rows["social_security_retirement__2024"].tolist() - ) == {22_000.0} - assert set( - synthetic_rows["employment_income_before_lsr__2024"].tolist() - ) == {50_000.0} + assert set(synthetic_rows["social_security_retirement__2024"].tolist()) == { + 22_000.0 + } + assert set(synthetic_rows["employment_income_before_lsr__2024"].tolist()) == { + 50_000.0 + } def test_support_augmentation_appends_mixed_age_household(): @@ -333,9 +335,11 @@ def test_support_augmentation_appends_mixed_age_household(): ) assert report["base_household_count"] == 3 assert report["augmented_household_count"] == 4 - synthetic_household_ids = set( - augmented_df["household_id__2024"].unique() - ) - {1, 2, 3} + synthetic_household_ids = set(augmented_df["household_id__2024"].unique()) - { + 1, + 2, + 3, + } assert len(synthetic_household_ids) == 1 synthetic_rows = augmented_df[ augmented_df["household_id__2024"].isin(synthetic_household_ids) @@ -459,9 +463,7 @@ def accessor(year: int): return SimpleNamespace( gov=SimpleNamespace( irs=SimpleNamespace( - payroll=SimpleNamespace( - social_security=SimpleNamespace(cap=cap) - ) + payroll=SimpleNamespace(social_security=SimpleNamespace(cap=cap)) ) ) ) @@ -918,9 +920,7 @@ def test_entropy_calibration_uses_lp_exact_fallback_even_before_approximate_wind monkeypatch.setattr( calibration_module, "calibrate_entropy", - lambda *args, **kwargs: (_ for _ in ()).throw( - RuntimeError("entropy stalled") - ), + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("entropy stalled")), ) monkeypatch.setattr( calibration_module, @@ -968,9 +968,7 @@ def test_nonnegative_feasibility_diagnostic_distinguishes_feasible_and_infeasibl ] ) infeasible_targets = np.array([1.0, 1.0, 3.0]) - infeasible = assess_nonnegative_feasibility( - infeasible_A, infeasible_targets - ) + infeasible = assess_nonnegative_feasibility(infeasible_A, infeasible_targets) assert infeasible["success"] is True assert infeasible["best_case_max_pct_error"] > 10.0 @@ -1163,11 +1161,12 @@ def test_manifest_persists_support_augmentation_metadata(tmp_path): metadata = json.loads(metadata_path.read_text(encoding="utf-8")) manifest = json.loads(manifest_path.read_text(encoding="utf-8")) assert metadata["support_augmentation"]["name"] == "donor-backed-synthetic-v1" - assert metadata["support_augmentation"]["report_file"] == "support_augmentation_report.json" assert ( - manifest["support_augmentation"]["report_summary"][ - "augmented_household_count" - ] + metadata["support_augmentation"]["report_file"] + == "support_augmentation_report.json" + ) + assert ( + manifest["support_augmentation"]["report_summary"]["augmented_household_count"] == 41326 )