From f68832e694c5bdc029005fbaead90ad3e658135f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 01:33:48 -0400 Subject: [PATCH 01/76] Document Social Security benefit lever implementation plan --- changelog.d/ss-benefit-levers.added.md | 1 + docs/social-security-benefit-levers.md | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+) create mode 100644 changelog.d/ss-benefit-levers.added.md create mode 100644 docs/social-security-benefit-levers.md diff --git a/changelog.d/ss-benefit-levers.added.md b/changelog.d/ss-benefit-levers.added.md new file mode 100644 index 00000000000..3720cded084 --- /dev/null +++ b/changelog.d/ss-benefit-levers.added.md @@ -0,0 +1 @@ +Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. diff --git a/docs/social-security-benefit-levers.md b/docs/social-security-benefit-levers.md new file mode 100644 index 00000000000..bd56516c0eb --- /dev/null +++ b/docs/social-security-benefit-levers.md @@ -0,0 +1,18 @@ +# Social Security benefit levers + +The benefit lever interface separates reported annual benefits from public +benefit outputs. The implementation plan is: + +1. Preserve the four reported benefit components as pure inputs, with the + existing effective uprating. Migrate legacy constructor inputs by period, + retaining explicitly supplied final overrides. +2. Apply independent across-the-board and supplied payable-share factors, and + one joint payment-year COLA factor with an initially inactive withholding + term. Keep the default factors equal to one. +3. Audit stored scoring-data columns and the artifact identity before scoring. +4. Verify current-law equality, precedence, periods, branches, export/reload, + factor properties, and a capped direct PE-US microsimulation. + +The engine contract does not certify API or application field exposure. +Calculated benefit source modes are reserved for subsequent implementations. +No package-specific values are part of this interface. From ad80c62ad071b7257279fd333c47f4546310544c Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 01:58:36 -0400 Subject: [PATCH 02/76] Add audited Social Security COLA and bend-point histories --- .../gov/ssa/cola_history/effective_month.yaml | 15 +++ .../ssa/cola_history/first_payment_month.yaml | 15 +++ .../gov/ssa/cola_history/pia_rounding.yaml | 21 ++++ .../parameters/gov/ssa/cola_history/rate.yaml | 70 +++++++++++ .../social_security/pia/formula_factors.yaml | 116 ++++++++++++++++++ .../tests/core/test_ssa_cola_history.py | 90 ++++++++++++++ policyengine_us/tools/ssa_cola_history.py | 45 +++++++ 7 files changed, 372 insertions(+) create mode 100644 policyengine_us/parameters/gov/ssa/cola_history/effective_month.yaml create mode 100644 policyengine_us/parameters/gov/ssa/cola_history/first_payment_month.yaml create mode 100644 policyengine_us/parameters/gov/ssa/cola_history/pia_rounding.yaml create mode 100644 policyengine_us/parameters/gov/ssa/cola_history/rate.yaml create mode 100644 policyengine_us/tests/core/test_ssa_cola_history.py create mode 100644 policyengine_us/tools/ssa_cola_history.py diff --git a/policyengine_us/parameters/gov/ssa/cola_history/effective_month.yaml b/policyengine_us/parameters/gov/ssa/cola_history/effective_month.yaml new file mode 100644 index 00000000000..b41cf18dc3b --- /dev/null +++ b/policyengine_us/parameters/gov/ssa/cola_history/effective_month.yaml @@ -0,0 +1,15 @@ +description: Social Security applies a historical annual COLA in this calendar month of its effective year. +values: + 1975-06-01: 6 + 1983-12-01: 12 +metadata: + unit: int + label: Social Security historical COLA effective month + economy: false + reference: + - title: Social Security Cost-Of-Living Adjustments + href: https://www.ssa.gov/OACT/COLA/colaseries.html +documentation: | + Read this month at an event date in cola_history.rate. June-effective + events apply in 1975-1982, and December-effective events apply from 1983. + No COLA occurred in December 1982 or June 1983. diff --git a/policyengine_us/parameters/gov/ssa/cola_history/first_payment_month.yaml b/policyengine_us/parameters/gov/ssa/cola_history/first_payment_month.yaml new file mode 100644 index 00000000000..605255718b5 --- /dev/null +++ b/policyengine_us/parameters/gov/ssa/cola_history/first_payment_month.yaml @@ -0,0 +1,15 @@ +description: Social Security first pays a historical annual COLA in this calendar month. +values: + 1975-06-01: 7 + 1983-12-01: 1 +metadata: + unit: int + label: Social Security historical COLA first payment month + economy: false + reference: + - title: Social Security COLA information, history of automatic adjustments + href: https://www.ssa.gov/news/en/cola/ +documentation: | + Read this month at an event date in cola_history.rate. For a June-effective + event, the first payment is in July of the same year. For a December-effective + event, the first payment is in January of the next year. diff --git a/policyengine_us/parameters/gov/ssa/cola_history/pia_rounding.yaml b/policyengine_us/parameters/gov/ssa/cola_history/pia_rounding.yaml new file mode 100644 index 00000000000..447daa26acf --- /dev/null +++ b/policyengine_us/parameters/gov/ssa/cola_history/pia_rounding.yaml @@ -0,0 +1,21 @@ +description: Social Security applies this historical rounding rule to the PIA after a COLA. +values: + 1975-06-01: 1 + 1982-06-01: 2 +metadata: + unit: int + label: Social Security historical COLA PIA rounding rule + economy: false + reference: + - title: POMS RS 00601.020 A, rounding of benefits + href: https://secure.ssa.gov/poms.nsf/lnx/0300601020 + - title: POMS RS 00605.021 C, computation example with historical COLAs + href: https://secure.ssa.gov/poms.nsf/lnx/0300605021 +documentation: | + Code 1 means round up to the next dime if not already a multiple of a dime. + Code 2 means truncate to the next lower dime. The June 1982 COLA uses code 2, + even where the initial PIA for a pre-1983 benchmark year rounded upward. + These codes describe COLA PIA rounding, not initial-PIA benchmark rounding. + A computation must reject any history containing a code it does not implement. + Reported-source proportional factors approximate changes in annual survey + amounts and do not reconstruct unidentified monthly benefit rounding. diff --git a/policyengine_us/parameters/gov/ssa/cola_history/rate.yaml b/policyengine_us/parameters/gov/ssa/cola_history/rate.yaml new file mode 100644 index 00000000000..9e26b7c528b --- /dev/null +++ b/policyengine_us/parameters/gov/ssa/cola_history/rate.yaml @@ -0,0 +1,70 @@ +description: Social Security applies these historical statutory cost-of-living adjustment rates, dated by their effective benefit month. +values: + 1975-06-01: 0.080 + 1976-06-01: 0.064 + 1977-06-01: 0.059 + 1978-06-01: 0.065 + 1979-06-01: 0.099 + 1980-06-01: 0.143 + 1981-06-01: 0.112 + 1982-06-01: 0.074 + 1983-12-01: 0.035 + 1984-12-01: 0.035 + 1985-12-01: 0.031 + 1986-12-01: 0.013 + 1987-12-01: 0.042 + 1988-12-01: 0.040 + 1989-12-01: 0.047 + 1990-12-01: 0.054 + 1991-12-01: 0.037 + 1992-12-01: 0.030 + 1993-12-01: 0.026 + 1994-12-01: 0.028 + 1995-12-01: 0.026 + 1996-12-01: 0.029 + 1997-12-01: 0.021 + 1998-12-01: 0.013 + 1999-12-01: 0.025 # Corrected from 0.024 by Public Law 106-554. + 2000-12-01: 0.035 + 2001-12-01: 0.026 + 2002-12-01: 0.014 + 2003-12-01: 0.021 + 2004-12-01: 0.027 + 2005-12-01: 0.041 + 2006-12-01: 0.033 + 2007-12-01: 0.023 + 2008-12-01: 0.058 + 2009-12-01: 0.000 + 2010-12-01: 0.000 + 2011-12-01: 0.036 + 2012-12-01: 0.017 + 2013-12-01: 0.015 + 2014-12-01: 0.017 + 2015-12-01: 0.000 + 2016-12-01: 0.003 + 2017-12-01: 0.020 + 2018-12-01: 0.028 + 2019-12-01: 0.016 + 2020-12-01: 0.013 + 2021-12-01: 0.059 + 2022-12-01: 0.087 + 2023-12-01: 0.032 + 2024-12-01: 0.025 + 2025-12-01: 0.028 +metadata: + unit: /1 + label: Social Security historical statutory COLA rate + economy: false + historical_start: 1975-06-01 + historical_end: 2025-12-01 + reference: + - title: Social Security Cost-Of-Living Adjustments + href: https://www.ssa.gov/OACT/COLA/colaseries.html +documentation: | + Dates are effective benefit months, not payment months. The June 1975-1982 + adjustments were first paid in July of the same year; adjustments from + December 1983 were first paid the following January. There was no December + 1982 or June 1983 adjustment. Explicit zero entries retain the 2009, 2010, + and 2015 determinations. December 1999 includes the statutory correction to + 2.5 percent. This audited series ends with December 2025, first paid January + 2026; later payment-year rates use the existing gov.ssa.uprating schedule. diff --git a/policyengine_us/parameters/gov/ssa/social_security/pia/formula_factors.yaml b/policyengine_us/parameters/gov/ssa/social_security/pia/formula_factors.yaml index 826166fbfdf..54e2b0ad0d6 100644 --- a/policyengine_us/parameters/gov/ssa/social_security/pia/formula_factors.yaml +++ b/policyengine_us/parameters/gov/ssa/social_security/pia/formula_factors.yaml @@ -9,9 +9,53 @@ metadata: href: https://www.law.cornell.edu/uscode/text/42/415#a_1_A - title: PIA Formula href: https://www.ssa.gov/OACT/COLA/piaformula.html + - title: Social Security benefit formula bend point history + href: https://www.ssa.gov/OACT/COLA/bendpoints.html +documentation: | + Each dated bend point is the amount for that initial eligibility year. + Historical values begin in 1979, when this wage-indexed formula began. + Their presence does not make + older computation methods supported. Calculated routes using modern initial + PIA truncation must reject benchmark years before 1983. brackets: - threshold: values: + 1979-01-01: 0 + 1980-01-01: 0 + 1981-01-01: 0 + 1982-01-01: 0 + 1983-01-01: 0 + 1984-01-01: 0 + 1985-01-01: 0 + 1986-01-01: 0 + 1987-01-01: 0 + 1988-01-01: 0 + 1989-01-01: 0 + 1990-01-01: 0 + 1991-01-01: 0 + 1992-01-01: 0 + 1993-01-01: 0 + 1994-01-01: 0 + 1995-01-01: 0 + 1996-01-01: 0 + 1997-01-01: 0 + 1998-01-01: 0 + 1999-01-01: 0 + 2000-01-01: 0 + 2001-01-01: 0 + 2002-01-01: 0 + 2003-01-01: 0 + 2004-01-01: 0 + 2005-01-01: 0 + 2006-01-01: 0 + 2007-01-01: 0 + 2008-01-01: 0 + 2009-01-01: 0 + 2010-01-01: 0 + 2011-01-01: 0 + 2012-01-01: 0 + 2013-01-01: 0 + 2014-01-01: 0 2015-01-01: 0 2016-01-01: 0 2017-01-01: 0 @@ -28,6 +72,42 @@ brackets: 1979-01-01: 0.90 - threshold: values: + 1979-01-01: 180 + 1980-01-01: 194 + 1981-01-01: 211 + 1982-01-01: 230 + 1983-01-01: 254 + 1984-01-01: 267 + 1985-01-01: 280 + 1986-01-01: 297 + 1987-01-01: 310 + 1988-01-01: 319 + 1989-01-01: 339 + 1990-01-01: 356 + 1991-01-01: 370 + 1992-01-01: 387 + 1993-01-01: 401 + 1994-01-01: 422 + 1995-01-01: 426 + 1996-01-01: 437 + 1997-01-01: 455 + 1998-01-01: 477 + 1999-01-01: 505 + 2000-01-01: 531 + 2001-01-01: 561 + 2002-01-01: 592 + 2003-01-01: 606 + 2004-01-01: 612 + 2005-01-01: 627 + 2006-01-01: 656 + 2007-01-01: 680 + 2008-01-01: 711 + 2009-01-01: 744 + 2010-01-01: 761 + 2011-01-01: 749 + 2012-01-01: 767 + 2013-01-01: 791 + 2014-01-01: 816 2015-01-01: 826 2016-01-01: 856 2017-01-01: 885 @@ -44,6 +124,42 @@ brackets: 1979-01-01: 0.32 - threshold: values: + 1979-01-01: 1_085 + 1980-01-01: 1_171 + 1981-01-01: 1_274 + 1982-01-01: 1_388 + 1983-01-01: 1_528 + 1984-01-01: 1_612 + 1985-01-01: 1_691 + 1986-01-01: 1_790 + 1987-01-01: 1_866 + 1988-01-01: 1_922 + 1989-01-01: 2_044 + 1990-01-01: 2_145 + 1991-01-01: 2_230 + 1992-01-01: 2_333 + 1993-01-01: 2_420 + 1994-01-01: 2_545 + 1995-01-01: 2_567 + 1996-01-01: 2_635 + 1997-01-01: 2_741 + 1998-01-01: 2_875 + 1999-01-01: 3_043 + 2000-01-01: 3_202 + 2001-01-01: 3_381 + 2002-01-01: 3_567 + 2003-01-01: 3_653 + 2004-01-01: 3_689 + 2005-01-01: 3_779 + 2006-01-01: 3_955 + 2007-01-01: 4_100 + 2008-01-01: 4_288 + 2009-01-01: 4_483 + 2010-01-01: 4_586 + 2011-01-01: 4_517 + 2012-01-01: 4_624 + 2013-01-01: 4_768 + 2014-01-01: 4_917 2015-01-01: 4_980 2016-01-01: 5_157 2017-01-01: 5_336 diff --git a/policyengine_us/tests/core/test_ssa_cola_history.py b/policyengine_us/tests/core/test_ssa_cola_history.py new file mode 100644 index 00000000000..9fccfdab66b --- /dev/null +++ b/policyengine_us/tests/core/test_ssa_cola_history.py @@ -0,0 +1,90 @@ +"""Historical statutory rounding arithmetic; no model construction required.""" + +from decimal import Decimal +from pathlib import Path + +import pytest +from policyengine_core.parameters import ParameterNode +from policyengine_core.parameters.helpers import load_parameter_file + +from policyengine_us.tools.ssa_cola_history import ( + apply_cola_to_pia, + validate_cola_rounding_rule, +) + + +def test_history_preserves_effective_dates_payment_timing_and_correction(): + history_path = ( + Path(__file__).resolve().parents[2] / "parameters/gov/ssa/cola_history" + ) + history = ParameterNode("cola_history", directory_path=str(history_path)) + events = {event.instant_str: event.value for event in history.rate.values_list} + + assert len(events) == 51 + assert min(events) == "1975-06-01" + assert max(events) == "2025-12-01" + assert "1982-12-01" not in events + assert "1983-06-01" not in events + assert history("1982-06-01").effective_month == 6 + assert history("1982-06-01").first_payment_month == 7 + assert history("1983-12-01").effective_month == 12 + assert history("1983-12-01").first_payment_month == 1 + assert events["1999-12-01"] == 0.025 + assert events["2009-12-01"] == events["2010-12-01"] == 0 + assert history("1981-06-01").pia_rounding == 1 + assert history("1982-06-01").pia_rounding == 2 + for date in events: + validate_cola_rounding_rule(history(date).pia_rounding) + + +def test_historical_bend_points_include_initial_and_declining_wage_years(): + bendpoint_path = ( + Path(__file__).resolve().parents[2] + / "parameters/gov/ssa/social_security/pia/formula_factors.yaml" + ) + scale = load_parameter_file(str(bendpoint_path), "pia_formula_factors") + assert scale("1979-01-01").thresholds == [0, 180, 1085] + assert scale("1983-01-01").thresholds == [0, 254, 1528] + assert scale("2010-01-01").thresholds == [0, 761, 4586] + assert scale("2011-01-01").thresholds == [0, 749, 4517] + assert scale("2014-01-01").thresholds == [0, 816, 4917] + + +def test_poms_historical_cola_rounding_sequence(): + # POMS RS 00605.021 C: June 1979, 1980, 1981 round up; June 1982 down. + pia = Decimal("322.70") + for rate, rule, expected in ( + ("0.099", 1, "354.70"), + ("0.143", 1, "405.50"), + ("0.112", 1, "451.00"), + ("0.074", 2, "484.30"), + ): + pia = apply_cola_to_pia(pia, rate, rule) + assert pia == Decimal(expected) + + +def test_ssa_modern_cola_rounding_example(): + # SSA's application-of-COLA example: 2,197.50 increased by 3.6 percent. + assert apply_cola_to_pia("2197.50", "0.036", 2) == Decimal("2276.60") + + +@pytest.mark.parametrize("rule", [1, 2]) +def test_exact_dime_is_unchanged(rule): + assert apply_cola_to_pia("100.10", 0, rule) == Decimal("100.1") + assert apply_cola_to_pia(0, "0.028", rule) == Decimal("0.0") + + +@pytest.mark.parametrize("rule", [0, 3, None, "down_to_dime", True]) +def test_unknown_historical_rounding_rule_is_rejected(rule): + with pytest.raises(ValueError, match="Unsupported historical COLA rounding rule"): + validate_cola_rounding_rule(rule) + with pytest.raises(ValueError, match="Unsupported historical COLA rounding rule"): + apply_cola_to_pia(100, "0.028", rule) + + +@pytest.mark.parametrize( + "pia, rate", [(-1, 0), (100, -0.01), ("NaN", 0), (100, "Infinity")] +) +def test_invalid_cola_inputs_are_rejected(pia, rate): + with pytest.raises(ValueError, match="finite and nonnegative"): + apply_cola_to_pia(pia, rate, 2) diff --git a/policyengine_us/tools/ssa_cola_history.py b/policyengine_us/tools/ssa_cola_history.py new file mode 100644 index 00000000000..c104b3caf62 --- /dev/null +++ b/policyengine_us/tools/ssa_cola_history.py @@ -0,0 +1,45 @@ +"""Pure historical COLA rounding, separate from reported-source factors. + +The initial-PIA benchmark-year rule is a different operation. This module +implements the rounding of an existing PIA after each historical COLA. +""" + +from decimal import Decimal, ROUND_CEILING, ROUND_DOWN + + +COLA_ROUND_UP_TO_DIME = 1 +COLA_ROUND_DOWN_TO_DIME = 2 + + +def validate_cola_rounding_rule(rounding_rule): + """Reject historical rules that have no implemented rounding operation. + + Codes come from ``gov.ssa.cola_history.pia_rounding``. Validate every + exposed event before using a cohort's history, including baseline events + which cancel from a proportional reform factor. + """ + if isinstance(rounding_rule, bool) or rounding_rule not in ( + COLA_ROUND_UP_TO_DIME, + COLA_ROUND_DOWN_TO_DIME, + ): + raise ValueError( + f"Unsupported historical COLA rounding rule: {rounding_rule!r}" + ) + return rounding_rule + + +def apply_cola_to_pia(pia, rate, rounding_rule): + """Apply one COLA to a monthly PIA and return an exactly rounded Decimal. + + POMS RS 00601.020 A specifies upward dime rounding before June 1982 and + downward dime rounding from June 1982. POMS RS 00605.021 C gives a worked + 1979-1982 sequence using both rules. Conversion through decimal strings + avoids importing binary floating-point error into a statutory dime boundary. + """ + rule = validate_cola_rounding_rule(rounding_rule) + pia = Decimal(str(pia)) + rate = Decimal(str(rate)) + if not pia.is_finite() or not rate.is_finite() or pia < 0 or rate < 0: + raise ValueError("PIA and statutory COLA rate must be finite and nonnegative") + rounding = ROUND_CEILING if rule == COLA_ROUND_UP_TO_DIME else ROUND_DOWN + return (pia * (1 + rate)).quantize(Decimal("0.1"), rounding=rounding) From 6758e661e72c2beeb2971fceda2d5460149c275a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 02:17:50 -0400 Subject: [PATCH 03/76] Audit stored Social Security scoring sources before applying levers --- .../social-security-scoring-preflight.md | 41 +++ .../test_social_security_preflight.py | 202 +++++++++++++++ .../tools/social_security_preflight.py | 245 ++++++++++++++++++ 3 files changed, 488 insertions(+) create mode 100644 docs/usage/social-security-scoring-preflight.md create mode 100644 policyengine_us/tests/code_health/test_social_security_preflight.py create mode 100644 policyengine_us/tools/social_security_preflight.py diff --git a/docs/usage/social-security-scoring-preflight.md b/docs/usage/social-security-scoring-preflight.md new file mode 100644 index 00000000000..e949b622c26 --- /dev/null +++ b/docs/usage/social-security-scoring-preflight.md @@ -0,0 +1,41 @@ +# Social Security scoring data preflight + +Audit the actual stored dataset before scoring Social Security reform levers: + +```bash +python -m policyengine_us.tools.social_security_preflight /path/to/populace_us_2024.h5 > receipt.json +``` + +The defaults identify the certified 2024 artifact +`hf://datasets/policyengine/populace-us/populace_us_2024.h5@populace-us-2024-spm-20260909` +and require SHA256 +`6496cc4393d4d3c6574f76eca231de5898c803b9067645591fd5c4d3e65aee84`. +An alternative artifact requires its own `--artifact`, `--year` and `--sha256`. +The preflight reads the H5's stored year and entity column metadata, without +calculating model outputs or loading the entity rows. + +Each of retirement, disability, survivors and dependents must have exactly one +Person source: its original public column name or its `_reported` input name. +The receipt records the source mapping that the PE-US constructor must implement. +Two names for one component would retain a final public override and are rejected. +Stored aggregate `social_security`, OASDI payroll-tax outputs, federal NIIT +liabilities/revenues and TOB revenue outputs are also rejected: these would bypass +the formulas being scored. This population-data contract permits household +situations to continue using explicit final benefit overrides. + +The JSON receipt records the artifact, hash, base year, stored H5 keys, complete +column inventory by entity, row counts, source mapping and any declared final +overrides. The Python function also accepts `intentional_final_overrides` so a +scoring caller can record its other overrides; final benefit and affected tax +overrides are rejected. The receipt applies to those exact bytes and that +declared override set. Constructor integration tests separately verify that each +legacy component source is migrated once. + +Use `extend_single_year_dataset` explicitly through the scoring horizon (2035 by +default; change the recorded horizon with `--projection-end-year`). Its projected +rows are independent annual cross-sections. Reconstruct retirement eligibility +within each scoring year as `payment_year - age + 62` when authoritative inputs +are unavailable. Thus a row aged 70 implies 2018 eligibility in a 2026 score and +2027 eligibility in a 2035 score. Integer ages leave a one-year birthday ambiguity. +Do not carry inferred eligibility, claiming, entitlement or prior-income histories +between projected rows; report the relevant proxy sensitivities with the score. diff --git a/policyengine_us/tests/code_health/test_social_security_preflight.py b/policyengine_us/tests/code_health/test_social_security_preflight.py new file mode 100644 index 00000000000..7944574f04d --- /dev/null +++ b/policyengine_us/tests/code_health/test_social_security_preflight.py @@ -0,0 +1,202 @@ +"""Stored-data audits must catch inputs that bypass benefit reform formulas.""" + +from hashlib import sha256 + +import pandas as pd +import pytest + +from policyengine_us.tools.social_security_preflight import ( + COMPONENTS, + preflight_social_security_dataset, +) + + +def make_dataset( + tmp_path, + *, + person_extra=None, + tax_unit_extra=None, + source_style="legacy", + year=2024, + storage_format="table", +): + person = {"person_id": [1], "age": [70]} + for public_name in COMPONENTS: + name = public_name if source_style == "legacy" else f"{public_name}_reported" + person[name] = [1_200] + person.update(person_extra or {}) + tax_unit = {"tax_unit_id": [1]} + tax_unit.update(tax_unit_extra or {}) + path = tmp_path / "synthetic.h5" + with pd.HDFStore(path, mode="w") as store: + store.put("person", pd.DataFrame(person), format=storage_format) + store.put( + "household", pd.DataFrame({"household_id": [1]}), format=storage_format + ) + store.put("tax_unit", pd.DataFrame(tax_unit), format=storage_format) + store.put("_time_period", pd.Series([year]), format=storage_format) + return path + + +def audit(path, **kwargs): + return preflight_social_security_dataset( + path, + artifact="synthetic://social-security-preflight", + expected_sha256=sha256(path.read_bytes()).hexdigest(), + **kwargs, + ) + + +@pytest.mark.parametrize("storage_format", ["table", "fixed"]) +def test_receipt_records_actual_stored_sources(tmp_path, storage_format): + path = make_dataset(tmp_path, storage_format=storage_format) + digest = sha256(path.read_bytes()).hexdigest() + receipt = audit(path, intentional_final_overrides=["employment_income"]) + + assert receipt["artifact"] == "synthetic://social-security-preflight" + assert receipt["sha256"] == digest + assert receipt["base_year"] == 2024 + assert receipt["stored_h5_keys"] == [ + "/_time_period", + "/household", + "/person", + "/tax_unit", + ] + assert receipt["stored_column_count"] == 8 + assert receipt["stored_columns"] == sorted( + ["person_id", "household_id", "tax_unit_id", "age", *COMPONENTS] + ) + assert receipt["row_counts"] == {"person": 1, "household": 1, "tax_unit": 1} + assert receipt["source_mapping"] == { + public_name: f"{public_name}_reported" for public_name in COMPONENTS + } + assert receipt["intentional_final_overrides"] == ["employment_income"] + assert receipt["projection"]["method"] == "extend_single_year_dataset" + assert receipt["projection"]["end_year"] == 2035 + assert receipt["projection"]["carry_inferred_history_across_years"] is False + assert sha256(path.read_bytes()).hexdigest() == digest + + +def test_reported_columns_are_not_migrated_again(tmp_path): + receipt = audit(make_dataset(tmp_path, source_style="reported")) + assert receipt["source_mapping"] == { + f"{public_name}_reported": f"{public_name}_reported" + for public_name in COMPONENTS + } + + +@pytest.mark.parametrize( + "column", + [ + "social_security", + "social_security_across_the_board_factor", + "social_security_payable_factor", + "social_security_retirement_base", + "social_security_retirement_joint_cola_factor", + "social_security_retirement_pia_ratio_factor", + "social_security_dependents_payable_factor", + "employee_social_security_tax", + "employer_social_security_tax", + "self_employment_social_security_tax", + "employer_total_social_security_tax", + "net_investment_income_tax", + "niit_liability", + "niit_revenue_oasdi", + "oasdi_niit_revenue", + "tob_revenue_oasdi", + "tob_revenue_medicare_hi", + "tob_revenue_total", + ], +) +def test_stored_outputs_rejected_even_when_zero(tmp_path, column): + path = make_dataset(tmp_path, tax_unit_extra={column: [0]}) + with pytest.raises(ValueError, match="bypass reform formulas"): + audit(path) + + +def test_both_component_names_would_preserve_a_bypassing_override(tmp_path): + path = make_dataset( + tmp_path, person_extra={"social_security_retirement_reported": [0]} + ) + with pytest.raises(ValueError, match="exactly one Person source"): + audit(path) + + +def test_missing_component_cannot_be_invented_from_aggregate(tmp_path): + path = make_dataset(tmp_path) + with pd.HDFStore(path, mode="a") as store: + person = store["person"].drop(columns=["social_security_disability"]) + store.put("person", person, format="table") + with pytest.raises(ValueError, match="exactly one Person source"): + audit(path) + + +@pytest.mark.parametrize( + "name", + [ + "social_security", + "social_security_retirement", + "tob_revenue_oasdi", + "social_security_payable_factor", + ], +) +def test_population_final_overrides_rejected(tmp_path, name): + path = make_dataset(tmp_path) + with pytest.raises(ValueError, match="cannot use final benefit or tax overrides"): + audit(path, intentional_final_overrides=[name]) + + +def test_hash_must_match_the_selected_artifact(tmp_path): + path = make_dataset(tmp_path) + with pytest.raises(ValueError, match="SHA256 mismatch"): + preflight_social_security_dataset( + path, + artifact="synthetic://social-security-preflight", + expected_sha256="0" * 64, + ) + + +def test_certified_identity_cannot_be_attached_to_alternative_bytes(tmp_path): + path = make_dataset(tmp_path) + with pytest.raises(ValueError, match="pins cannot be changed"): + preflight_social_security_dataset( + path, expected_sha256=sha256(path.read_bytes()).hexdigest() + ) + + +def test_base_year_comes_from_h5_not_the_filename(tmp_path): + path = make_dataset(tmp_path, year=2023) + with pytest.raises(ValueError, match="base year mismatch"): + audit(path) + + +def test_base_year_must_be_stored(tmp_path): + path = make_dataset(tmp_path) + with pd.HDFStore(path, mode="a") as store: + store.remove("_time_period") + with pytest.raises(ValueError, match="must store its base year"): + audit(path) + + +def test_columns_cannot_be_duplicated_across_entities(tmp_path): + path = make_dataset(tmp_path, tax_unit_extra={"social_security_retirement": [0]}) + with pytest.raises(ValueError, match="Columns appear more than once"): + audit(path) + + +def test_projection_horizon_must_include_base_year(tmp_path): + path = make_dataset(tmp_path) + with pytest.raises(ValueError, match="end year precedes"): + audit(path, projection_end_year=2023) + + +def test_table_inventory_does_not_read_entity_rows(tmp_path, monkeypatch): + path = make_dataset(tmp_path) + original_select = pd.HDFStore.select + + def select_metadata_only(self, key, *args, **kwargs): + assert key.strip("/") == "_time_period", "Entity rows must not be loaded." + return original_select(self, key, *args, **kwargs) + + monkeypatch.setattr(pd.HDFStore, "select", select_metadata_only) + assert audit(path)["stored_column_count"] == 8 diff --git a/policyengine_us/tools/social_security_preflight.py b/policyengine_us/tools/social_security_preflight.py new file mode 100644 index 00000000000..4a6f4f311f6 --- /dev/null +++ b/policyengine_us/tools/social_security_preflight.py @@ -0,0 +1,245 @@ +"""Audit stored sources before scoring Social Security benefit reforms. + +This reads HDF metadata, not calculated simulation outputs. A receipt applies +only to the exact bytes it hashes. It is deliberately a separate scoring +preflight: household situations may still supply final benefit overrides. +""" + +import argparse +from collections import Counter +from collections.abc import Iterable +from hashlib import sha256 +import json +from pathlib import Path + +import pandas as pd + + +CERTIFIED_ARTIFACT = ( + "hf://datasets/policyengine/populace-us/populace_us_2024.h5" + "@populace-us-2024-spm-20260909" +) +CERTIFIED_YEAR = 2024 +CERTIFIED_SHA256 = "6496cc4393d4d3c6574f76eca231de5898c803b9067645591fd5c4d3e65aee84" +COMPONENTS = ( + "social_security_retirement", + "social_security_disability", + "social_security_survivors", + "social_security_dependents", +) +CALCULATED_BENEFIT_LEAVES = frozenset( + { + "social_security", + "social_security_across_the_board_factor", + "social_security_payable_factor", + } + | { + f"{component}_{suffix}" + for component in COMPONENTS + for suffix in ( + "base", + "joint_cola_factor", + "pia_ratio_factor", + "payable_factor", + ) + } +) +CALCULATED_TAX_LEAVES = frozenset( + ( + "employee_social_security_tax", + "employer_social_security_tax", + "self_employment_social_security_tax", + "employer_total_social_security_tax", + "net_investment_income_tax", + "niit", + "oasdi_payroll_tax_revenue", + "oasdi_benefit_taxation_revenue", + "oasdi_niit_revenue", + "oasdi_income", + "oasdi_cost", + ) +) + + +def _is_calculated_tax_leaf(name: str) -> bool: + # Include later NIIT and TOB revenue variants without rejecting source + # inputs such as net investment income or a supplied NIIT income basis. + return ( + name in CALCULATED_TAX_LEAVES + or name.startswith("tob_revenue_") + or name.startswith("net_investment_income_tax_") + or name.startswith("niit_liability") + or name.startswith("niit_revenue") + or name.endswith("_niit_revenue") + ) + + +def _file_sha256(path: Path) -> str: + digest = sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def preflight_social_security_dataset( + file_path: str | Path, + *, + artifact: str = CERTIFIED_ARTIFACT, + expected_year: int = CERTIFIED_YEAR, + expected_sha256: str = CERTIFIED_SHA256, + projection_end_year: int = 2035, + intentional_final_overrides: Iterable[str] = (), +) -> dict: + """Validate a single-year stored H5 artifact and return a JSON receipt. + + Alternative artifacts must supply their own artifact URI, base year and + expected SHA256. Each component may be stored under its original name + (which the PE-US constructor migrates) or its pure ``_reported`` name. + Storing both is rejected because that preserves a final public override. + + Projection must use ``extend_single_year_dataset`` through the recorded + end year. Its rows are independent annual cross-sections: integer-age + retirement eligibility is inferred afresh as ``year - age + 62``, subject + to a one-year birthday ambiguity. This receipt does not turn projected + rows into a longitudinal entitlement or income history. + """ + path = Path(file_path) + if not artifact or not expected_sha256: + raise ValueError("An artifact identity and expected SHA256 are required.") + if artifact == CERTIFIED_ARTIFACT and ( + expected_year != CERTIFIED_YEAR or expected_sha256.lower() != CERTIFIED_SHA256 + ): + raise ValueError( + "The certified artifact's year and SHA256 pins cannot be changed; " + "identify an alternative artifact for an alternative audit." + ) + if projection_end_year < expected_year: + raise ValueError("The projection end year precedes the base year.") + actual_sha256 = _file_sha256(path) + if actual_sha256 != expected_sha256.lower(): + raise ValueError( + f"Dataset SHA256 mismatch: expected {expected_sha256}, got {actual_sha256}." + ) + + columns_by_entity = {} + row_counts = {} + entities = ("person", "household", "tax_unit", "spm_unit", "family", "marital_unit") + with pd.HDFStore(path, mode="r") as store: + keys = set(store.keys()) + if "/_time_period" not in keys: + raise ValueError("Dataset must store its base year in _time_period.") + stored_years = store["_time_period"].tolist() + if stored_years != [expected_year]: + raise ValueError( + f"Dataset base year mismatch: expected {expected_year}, " + f"got {stored_years}." + ) + unknown = keys - {"/_time_period"} - {f"/{entity}" for entity in entities} + if unknown: + raise ValueError(f"Unexpected stored H5 keys: {sorted(unknown)}.") + for entity in entities: + key = f"/{entity}" + if key not in keys: + if entity in ("person", "household", "tax_unit"): + raise ValueError(f"Dataset is missing required entity {entity}.") + continue + storer = store.get_storer(key) + if storer.is_table: + columns = list(storer.non_index_axes[0][1]) + row_counts[entity] = int(storer.nrows) + else: + # Fixed-format H5 stores column labels separately; slicing an + # empty frame reads the labels without materializing rows. + columns = list(store.select(key, start=0, stop=0).columns) + row_counts[entity] = int(storer.shape[0]) + columns_by_entity[entity] = sorted(columns) + + all_columns = [ + column for columns in columns_by_entity.values() for column in columns + ] + duplicate_columns = sorted( + name for name, count in Counter(all_columns).items() if count > 1 + ) + if duplicate_columns: + raise ValueError(f"Columns appear more than once: {duplicate_columns}.") + stored_columns = set(all_columns) + bypasses = sorted( + name + for name in stored_columns + if name in CALCULATED_BENEFIT_LEAVES or _is_calculated_tax_leaf(name) + ) + if bypasses: + raise ValueError( + f"Stored calculated outputs bypass reform formulas: {bypasses}." + ) + + source_mapping = {} + person_columns = set(columns_by_entity["person"]) + for public_name in COMPONENTS: + source_name = f"{public_name}_reported" + candidates = person_columns & {public_name, source_name} + if len(candidates) != 1: + raise ValueError( + f"Store exactly one Person source for {public_name}: " + f"{public_name} or {source_name}; found {sorted(candidates)}." + ) + stored_name = next(iter(candidates)) + source_mapping[stored_name] = source_name + + overrides = sorted(set(intentional_final_overrides)) + forbidden_overrides = sorted( + name + for name in overrides + if name in CALCULATED_BENEFIT_LEAVES + or name in COMPONENTS + or _is_calculated_tax_leaf(name) + ) + if forbidden_overrides: + raise ValueError( + "Population scoring cannot use final benefit or tax overrides: " + f"{forbidden_overrides}." + ) + + return { + "artifact": artifact, + "sha256": actual_sha256, + "base_year": expected_year, + "stored_h5_keys": sorted(keys), + "stored_columns": sorted(stored_columns), + "stored_column_count": len(stored_columns), + "columns_by_entity": columns_by_entity, + "row_counts": row_counts, + "source_mapping": source_mapping, + "intentional_final_overrides": overrides, + "projection": { + "method": "extend_single_year_dataset", + "end_year": projection_end_year, + "population": "independent annual cross-sections", + "retirement_eligibility_proxy": "payment_year - age + 62", + "integer_age_uncertainty_years": 1, + "carry_inferred_history_across_years": False, + }, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("file_path", type=Path) + parser.add_argument("--artifact", default=CERTIFIED_ARTIFACT) + parser.add_argument("--year", type=int, default=CERTIFIED_YEAR) + parser.add_argument("--sha256", default=CERTIFIED_SHA256) + parser.add_argument("--projection-end-year", type=int, default=2035) + args = parser.parse_args() + receipt = preflight_social_security_dataset( + args.file_path, + artifact=args.artifact, + expected_year=args.year, + expected_sha256=args.sha256, + projection_end_year=args.projection_end_year, + ) + print(json.dumps(receipt, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() From 1d17f32f8315b646a500a3022570e69bfcac1138 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 02:18:39 -0400 Subject: [PATCH 04/76] Separate reported Social Security inputs and compose benefit levers --- docs/social-security-benefit-levers.md | 147 ++++++- policyengine_us/data/economic_assumptions.py | 6 + .../benefit_levers/across_the_board/rate.yaml | 7 + .../across_the_board/start_year.yaml | 7 + .../benefit_levers/cola/alternative_rate.yaml | 7 + .../benefit_levers/cola/effective_month.yaml | 7 + .../benefit_levers/cola/effective_year.yaml | 7 + .../ssa/benefit_levers/cola/end_year.yaml | 7 + .../ssa/benefit_levers/cola/pp_change.yaml | 7 + .../ssa/benefit_levers/cola/start_year.yaml | 7 + .../cola/withholding_enabled.yaml | 7 + .../ssa/benefit_levers/payable/oasi_only.yaml | 7 + .../ssa/benefit_levers/payable/share.yaml | 7 + policyengine_us/system.py | 5 + .../core/test_social_security_current_law.py | 351 +++++++++++++++ .../tests/core/test_social_security_inputs.py | 305 +++++++++++++ .../test_social_security_lever_properties.py | 399 ++++++++++++++++++ ...security_dependents_joint_cola_factor.yaml | 57 +++ .../social_security_retirement.yaml | 127 ++++++ ...security_retirement_joint_cola_factor.yaml | 118 ++++++ .../tools/social_security_benefit_levers.py | 257 +++++++++++ .../tools/social_security_inputs.py | 152 +++++++ ...social_security_across_the_board_factor.py | 16 + .../gov/ssa/ss/social_security_dependents.py | 19 +- .../ssa/ss/social_security_dependents_base.py | 15 + ...ocial_security_dependents_cola_withheld.py | 10 + ...l_security_dependents_entitlement_month.py | 10 + ...al_security_dependents_entitlement_year.py | 10 + ...al_security_dependents_first_cola_month.py | 10 + ...ial_security_dependents_first_cola_year.py | 10 + .../ssa/ss/social_security_dependents_fund.py | 12 + ...l_security_dependents_joint_cola_factor.py | 17 + ...cial_security_dependents_payable_factor.py | 22 + ...al_security_dependents_pia_ratio_factor.py | 14 + .../ss/social_security_dependents_reported.py | 10 + .../ss/social_security_dependents_source.py | 13 + ...rity_dependents_worker_eligibility_year.py | 10 + .../social_security_dependents_worker_id.py | 10 + .../gov/ssa/ss/social_security_disability.py | 21 +- .../ssa/ss/social_security_disability_base.py | 15 + ...ocial_security_disability_cola_withheld.py | 10 + ...l_security_disability_entitlement_month.py | 10 + ...al_security_disability_entitlement_year.py | 10 + ...al_security_disability_first_cola_month.py | 10 + ...ial_security_disability_first_cola_year.py | 10 + ...l_security_disability_joint_cola_factor.py | 17 + ...cial_security_disability_payable_factor.py | 13 + ...al_security_disability_pia_ratio_factor.py | 14 + .../ss/social_security_disability_reported.py | 10 + .../ss/social_security_disability_source.py | 13 + ...rity_disability_worker_eligibility_year.py | 10 + .../social_security_disability_worker_id.py | 10 + .../ssa/ss/social_security_payable_factor.py | 15 + .../gov/ssa/ss/social_security_retirement.py | 19 +- .../ssa/ss/social_security_retirement_base.py | 15 + ...ocial_security_retirement_cola_withheld.py | 10 + ...l_security_retirement_entitlement_month.py | 10 + ...al_security_retirement_entitlement_year.py | 10 + ...al_security_retirement_first_cola_month.py | 10 + ...ial_security_retirement_first_cola_year.py | 10 + ...l_security_retirement_joint_cola_factor.py | 17 + ...cial_security_retirement_payable_factor.py | 13 + ...al_security_retirement_pia_ratio_factor.py | 14 + .../ss/social_security_retirement_reported.py | 10 + .../ss/social_security_retirement_source.py | 13 + ...rity_retirement_worker_eligibility_year.py | 10 + .../social_security_retirement_worker_id.py | 10 + .../gov/ssa/ss/social_security_survivors.py | 19 +- .../ssa/ss/social_security_survivors_base.py | 15 + ...social_security_survivors_cola_withheld.py | 10 + ...al_security_survivors_entitlement_month.py | 10 + ...ial_security_survivors_entitlement_year.py | 10 + ...ial_security_survivors_first_cola_month.py | 10 + ...cial_security_survivors_first_cola_year.py | 10 + ...al_security_survivors_joint_cola_factor.py | 17 + ...ocial_security_survivors_payable_factor.py | 13 + ...ial_security_survivors_pia_ratio_factor.py | 14 + .../ss/social_security_survivors_reported.py | 10 + .../ss/social_security_survivors_source.py | 13 + ...urity_survivors_worker_eligibility_year.py | 10 + .../ss/social_security_survivors_worker_id.py | 10 + ...social_security_worker_computation_year.py | 10 + 82 files changed, 2733 insertions(+), 16 deletions(-) create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/across_the_board/rate.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/across_the_board/start_year.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/alternative_rate.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/effective_month.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/effective_year.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/end_year.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/pp_change.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/start_year.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/withholding_enabled.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/payable/oasi_only.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/payable/share.yaml create mode 100644 policyengine_us/tests/core/test_social_security_current_law.py create mode 100644 policyengine_us/tests/core/test_social_security_inputs.py create mode 100644 policyengine_us/tests/core/test_social_security_lever_properties.py create mode 100644 policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_dependents_joint_cola_factor.yaml create mode 100644 policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement.yaml create mode 100644 policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_joint_cola_factor.yaml create mode 100644 policyengine_us/tools/social_security_benefit_levers.py create mode 100644 policyengine_us/tools/social_security_inputs.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_across_the_board_factor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_base.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_cola_withheld.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_entitlement_month.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_entitlement_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_first_cola_month.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_first_cola_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_fund.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_joint_cola_factor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_payable_factor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_factor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_reported.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_source.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_eligibility_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_id.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_base.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_cola_withheld.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_entitlement_month.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_entitlement_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_first_cola_month.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_first_cola_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_joint_cola_factor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_payable_factor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_factor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_reported.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_source.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_worker_eligibility_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_worker_id.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_payable_factor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_base.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_cola_withheld.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_entitlement_month.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_entitlement_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_first_cola_month.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_first_cola_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_joint_cola_factor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_payable_factor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_factor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_reported.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_source.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_worker_eligibility_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_worker_id.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_base.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_cola_withheld.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_entitlement_month.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_entitlement_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_first_cola_month.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_first_cola_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_joint_cola_factor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_payable_factor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_factor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_reported.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_source.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_eligibility_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_id.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_worker_computation_year.py diff --git a/docs/social-security-benefit-levers.md b/docs/social-security-benefit-levers.md index bd56516c0eb..a045fd076cf 100644 --- a/docs/social-security-benefit-levers.md +++ b/docs/social-security-benefit-levers.md @@ -1,18 +1,143 @@ # Social Security benefit levers The benefit lever interface separates reported annual benefits from public -benefit outputs. The implementation plan is: - -1. Preserve the four reported benefit components as pure inputs, with the - existing effective uprating. Migrate legacy constructor inputs by period, - retaining explicitly supplied final overrides. -2. Apply independent across-the-board and supplied payable-share factors, and - one joint payment-year COLA factor with an initially inactive withholding - term. Keep the default factors equal to one. -3. Audit stored scoring-data columns and the artifact identity before scoring. -4. Verify current-law equality, precedence, periods, branches, export/reload, - factor properties, and a capped direct PE-US microsimulation. +benefit outputs. Stable public formulas combine a selected reported source with +independent across-the-board and supplied payable-share factors, and one joint +payment-year COLA factor. All factors are one by default. A stored-data preflight +checks the scoring artifact's identity and columns before population scoring. The engine contract does not certify API or application field exposure. Calculated benefit source modes are reserved for subsequent implementations. No package-specific values are part of this interface. + +## Input and override contract + +The public components are `social_security_retirement`, +`social_security_disability`, `social_security_survivors`, and +`social_security_dependents`. Their `_reported` twins are pure annual Person +inputs, including zero, with the original effective per-capita SOI/CBO +uprating. Public formulas have no variable-level uprating. Dataset extension +still uprates original column names before constructor migration. + +Both PE-US constructors migrate a legacy public input to its reported twin +separately for each supplied period. When both names are supplied in the same +period, the reported amount remains the source and the public input remains a +final override. Later public `set_input` calls and aggregate `social_security` +inputs are final overrides. Core `SimulationBuilder` and YAML bypass constructor +migration: use `_reported` for lever-sensitive fixtures. A public override in +one year does not become the reported source for later years. + +Each component has a `_source` selector with default `reported`. Reserved +options `supplied_aime`, `supplied_pia`, `supplied_ratio_survivor`, and `history` +fail explicitly until their calculation routes are implemented. Selection never +depends on whether an amount is positive. The selected `_base` and the +`_pia_ratio_factor` (currently one) are separate from reported inputs. Future +calculated routes must bypass population PIA and COLA factors, apply the other +factors once, and perform their earnings test on the resulting payable amount. + +`policyengine_us.tools.social_security_inputs.export_source_inputs(simulation)` +exports all recorded pure-input periods and entity memberships as `sources`, +with supplied formula-variable values in a separate `overrides` payload. Cached +calculations are excluded. `reload_source_inputs(payload, ...)` restores sources +before final overrides. Core's legacy `extract_person` remains a one-period +interface. Neither interface settles the separate API/application metadata and +field-exposure contract. + +## Parameters and composition + +All lever parameters live under `gov.contrib.ssa.benefit_levers`: + +| Parameter | Default | Meaning | +| --- | --- | --- | +| `across_the_board.start_year` | 9999 | First affected payment year | +| `across_the_board.rate` | 0 | Proportional change, finite and at least -1 | +| `payable.share` | 1 | Externally supplied annual share in [0,1] | +| `payable.oasi_only` | false | OASI-only rather than combined OASDI | +| `cola.start_year` | 9999 | First affected January payment year | +| `cola.effective_year` | 0 | Optional first affected COLA effective year | +| `cola.effective_month` | 12 | Effective month when effective year is supplied | +| `cola.end_year` | 0 | Last affected payment year, or no end | +| `cola.pp_change` | 0 | Addition to the annual proportional COLA rate | +| `cola.alternative_rate` | -1 | Supplied pinned alternative annual rate; -1 selects baseline | +| `cola.withholding_enabled` | false | Enable component payment-year withholding terms | + +A positive `effective_year` takes precedence over `start_year`: December 2028 +means the first affected check is January 2029. Parameter values can vary by +payment year. Example modeling assumptions are CPI-E `pp_change = 0.0015` +(+0.15 percentage points), or chained CPI `pp_change = -0.0030` +(-0.30 percentage points), with an explicitly selected start. These are examples, +not observed index series or package presets. An alternative-rate path must be +pinned and recorded by its supplier; it is not generated from a household. + +For each exposed payment event, statutory baseline rate c and configured +addition delta give adjusted rate max(0,c+delta), without a second rounding of +the policy addition. The joint multiplier is `(1+adjusted)/(1+c)` unless that +event is withheld, in which case it is `1/(1+c)`. Each component's +`_cola_withheld` term defaults false; a subsequent means-test implementation +supplies the rules. Switching a policy off later preserves its earlier effect +on the benefit level. Public reported-route benefits multiply this joint factor, +the PIA ratio factor, the across-the-board factor, and the payable factor once. + +The payable supplier specifies an annual path beginning in the depletion year +and supplies a blended share for a depletion year beginning partway through the +year. The engine does not predict depletion, share paths, or within-year fund +cash flow. In OASI-only mode, disability is unaffected; a positive dependent +benefit requires an explicit `social_security_dependents_fund` of `oasi` or `di`. +Mixed or unidentified dependent streams require an external allocation/blended +scenario; the engine does not assign the entire stream to OASI silently. + +## COLA history and exposure + +Historical statutory events are under `gov.ssa.cola_history` through December +2025 (January 2026 payment), including June/July events before 1983, zero-rate +years and the corrected December 1999 2.5% rate. The current main January 2027 +SSA schedule already includes the actual Q3 2026 index. Starting with that +schedule entry, rate construction uses `gov.ssa.uprating`, comparing with the Q3 +index associated with the last effective nonzero COLA. Decimal half-up rounding +to 0.1 percentage point preserves the high-water base over zero-COLA years. +`cola_events(system.parameters, through_year)` in +`tools.social_security_benefit_levers` exports the exact baseline schedule for a +scoring receipt. Freeze this schedule together with the reform parameters. + +Historical PIA COLA rounding is implemented separately by +`tools.ssa_cola_history.apply_cola_to_pia`: round upward to a dime before June +1982 and downward thereafter; unknown rule codes fail explicitly. Reported-route +ratios validate exposed rules but do not reconstruct a rounded individual PIA +from an annual recipient amount. Initial-PIA computations/inverses before the +1983 benchmark remain outside this implementation. Bend-point history is now +available back to 1979 for future services. + +Per-component optional inputs are `_first_cola_year`/`_first_cola_month`, +`_worker_eligibility_year`, `_worker_id`, and +`_entitlement_year`/`_entitlement_month`. A first-COLA date takes precedence. +A worker link follows an explicit component basis on the linked worker, falling +back to the shared authoritative `social_security_worker_computation_year`. +A supplied recipient basis is preserved; conflicting positive bases or invalid +links fail. No disability or deceased-worker computation basis is inferred from +a linked worker's retirement age. + +Absent identified dates, own retirement uses age-62 eligibility (even if claiming +later); older survivors use an age-60 recipient proxy. Positive disability, +child/caregiver dependents and young survivors use a pre-start exposure proxy. +These proxies are assumptions, not observed histories. Supply current-year +entitlement/first-COLA dates as the lower-exposure sensitivity; compare adjacent +retirement eligibility years for integer-age birthday uncertainty and age-50 +survivor or linked-record alternatives where relevant. Report these sensitivity +envelopes separately from identified results. Without worker/share information, +auxiliary PIA inversion remains unidentified. + +Each projected cross-section infers eligibility anew from `year - age + 62`. +It does not carry cohorts or income histories across static population rows. +A supplied retirement entitlement/claim date does not erase pre-claim COLAs. +Pre-1983 retirement eligibility can expose a July payment in its eligibility +year; modern retirement eligibility first exposes January of the next year. + +Run the stored-data preflight described in +[the scoring-data guide](usage/social-security-scoring-preflight.md) before +population lever scoring. Its receipt records the artifact, bytes hash, stored +year, full stored columns, source mapping and projection convention. + +For a supplied alternative-index level path, convert its pinned levels (including +an explicit pre-start anchor) with `projected_cola_rates`, then supply the +resulting payment-year rates under `cola.alternative_rate`. Record the source, +anchor, level path and generated rate path with the scoring receipt. diff --git a/policyengine_us/data/economic_assumptions.py b/policyengine_us/data/economic_assumptions.py index 14347e5fb70..630978bc736 100644 --- a/policyengine_us/data/economic_assumptions.py +++ b/policyengine_us/data/economic_assumptions.py @@ -42,6 +42,12 @@ "self_employed_health_insurance_ald": DEFAULT_MICRODATA_UPRATING, "self_employed_pension_contribution_ald": DEFAULT_MICRODATA_UPRATING, "social_security": "calibration.gov.irs.soi.social_security", + # Dataset extension precedes constructor migration to the reported twins. + # Resolve these national totals through the existing per-capita uprater. + "social_security_retirement": "calibration.gov.irs.soi.social_security", + "social_security_disability": "calibration.gov.irs.soi.social_security", + "social_security_survivors": "calibration.gov.irs.soi.social_security", + "social_security_dependents": "calibration.gov.irs.soi.social_security", "s_corp_income": "calibration.gov.irs.soi.partnership_s_corp_income", "spm_unit_weight": "calibration.gov.census.populations.total", "state_and_local_sales_or_income_tax": DEFAULT_MICRODATA_UPRATING, diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/across_the_board/rate.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/across_the_board/rate.yaml new file mode 100644 index 00000000000..b70a3bd4b04 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/across_the_board/rate.yaml @@ -0,0 +1,7 @@ +description: "Proportional across-the-board benefit change, at least -1." +values: + 0000-01-01: 0 +metadata: + unit: /1 + label: Social Security across the board rate + economy: false diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/across_the_board/start_year.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/across_the_board/start_year.yaml new file mode 100644 index 00000000000..7a098179d2d --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/across_the_board/start_year.yaml @@ -0,0 +1,7 @@ +description: "First payment year of the across-the-board change." +values: + 0000-01-01: 9999 +metadata: + unit: year + label: Social Security across the board start year + economy: false diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/alternative_rate.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/alternative_rate.yaml new file mode 100644 index 00000000000..29a263f39c8 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/alternative_rate.yaml @@ -0,0 +1,7 @@ +description: "Optional externally pinned annual alternative COLA rate before pp_change, or -1 to use the baseline statutory rate." +values: + 0000-01-01: -1 +metadata: + unit: /1 + label: Social Security cola alternative rate + economy: false diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/effective_month.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/effective_month.yaml new file mode 100644 index 00000000000..dafde46fbb8 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/effective_month.yaml @@ -0,0 +1,7 @@ +description: "First affected COLA effective month when effective_year is supplied." +values: + 0000-01-01: 12 +metadata: + unit: month + label: Social Security cola effective month + economy: false diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/effective_year.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/effective_year.yaml new file mode 100644 index 00000000000..b61974d9d2a --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/effective_year.yaml @@ -0,0 +1,7 @@ +description: "First affected COLA effective calendar year, or zero to use start_year." +values: + 0000-01-01: 0 +metadata: + unit: year + label: Social Security cola effective year + economy: false diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/end_year.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/end_year.yaml new file mode 100644 index 00000000000..f1f86f725db --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/end_year.yaml @@ -0,0 +1,7 @@ +description: "Last affected payment calendar year, or zero for no end." +values: + 0000-01-01: 0 +metadata: + unit: year + label: Social Security cola end year + economy: false diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/pp_change.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/pp_change.yaml new file mode 100644 index 00000000000..f3b94f89f2e --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/pp_change.yaml @@ -0,0 +1,7 @@ +description: "Annual COLA proportional-rate addition; examples: CPI-E +0.15 percentage points = 0.0015, chained CPI -0.30 percentage points = -0.0030. Preserve this precision after statutory rounding." +values: + 0000-01-01: 0 +metadata: + unit: /1 + label: Social Security cola pp change + economy: false diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/start_year.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/start_year.yaml new file mode 100644 index 00000000000..ae728d81e44 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/start_year.yaml @@ -0,0 +1,7 @@ +description: "First affected January payment year. Used unless effective_year is supplied." +values: + 0000-01-01: 9999 +metadata: + unit: year + label: Social Security cola start year + economy: false diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/withholding_enabled.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/withholding_enabled.yaml new file mode 100644 index 00000000000..dd9e94d981e --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/withholding_enabled.yaml @@ -0,0 +1,7 @@ +description: "Enable the component payment-year withholding terms; subsequent means-test rules supply the decisions." +values: + 0000-01-01: false +metadata: + unit: bool + label: Social Security cola withholding enabled + economy: false diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/payable/oasi_only.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/payable/oasi_only.yaml new file mode 100644 index 00000000000..95fb0d2386c --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/payable/oasi_only.yaml @@ -0,0 +1,7 @@ +description: "Apply the supplied payable share only to OASI components rather than combined OASDI." +values: + 0000-01-01: false +metadata: + unit: bool + label: Social Security payable oasi only + economy: false diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/payable/share.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/payable/share.yaml new file mode 100644 index 00000000000..921b00d37e9 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/payable/share.yaml @@ -0,0 +1,7 @@ +description: "Supplied annual payable share in [0, 1]. The supplier blends any partial depletion year." +values: + 0000-01-01: 1 +metadata: + unit: /1 + label: Social Security payable share + economy: false diff --git a/policyengine_us/system.py b/policyengine_us/system.py index 1b8a070ec6b..25bf0c49b73 100644 --- a/policyengine_us/system.py +++ b/policyengine_us/system.py @@ -36,6 +36,7 @@ ) from .tools.default_uprating import add_default_uprating from .tools.period_branch import delete_input_values +from .tools.social_security_inputs import migrate_social_security_inputs from .tools.section_911 import validate_section_911_batch_inputs from .tools.per_capita_uprating import ( add_per_capita_parameters_for_parameter_uprating, @@ -259,6 +260,8 @@ def __init__(self, *args, **kwargs): if reform is not None: self.apply_reform(reform) + migrate_social_security_inputs(self) + # Labor supply responses employment_income = self.get_holder("employment_income") @@ -519,6 +522,8 @@ def __init__(self, *args, **kwargs): if reform is not None: self.apply_reform(reform) + migrate_social_security_inputs(self) + # Labor supply responses employment_income = self.get_holder("employment_income") diff --git a/policyengine_us/tests/core/test_social_security_current_law.py b/policyengine_us/tests/core/test_social_security_current_law.py new file mode 100644 index 00000000000..45e562e3095 --- /dev/null +++ b/policyengine_us/tests/core/test_social_security_current_law.py @@ -0,0 +1,351 @@ +"""Differential compatibility with the four original public benefit inputs. + +Python is required for constructor, Core builder, dataset, cache/branch and +export APIs that YAML cannot exercise. The reference restores the original +four pure input definitions and disables only the new constructor migration; +all other current-law policy and wrapper behavior is identical. One model +clone is shared read-only across the module. Every population has two people, +with no external dataset, and only the neutral-reform branches copy policy. +Local runtime/memory measurements are reported with the PR's test receipt. +""" + +from copy import deepcopy +from unittest.mock import patch + +import numpy as np +import pandas as pd +import pytest +from policyengine_core.periods import YEAR +from policyengine_core.reforms import Reform +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.variables import Variable + +from policyengine_us import Microsimulation, Simulation +from policyengine_us.system import system +from policyengine_us.tools.period_branch import drop_inherited_values +from policyengine_us.tools.social_security_inputs import ( + SOCIAL_SECURITY_COMPONENTS, + export_source_inputs, + reload_source_inputs, +) + + +# All six social_security* variables present before this change, including +# the payroll-tax base whose name shares the benefit prefix. +PREEXISTING_SOCIAL_SECURITY = ( + *SOCIAL_SECURITY_COMPONENTS, + "social_security", + "social_security_taxable_self_employment_income", +) +DOWNSTREAM = ( + "taxable_social_security", + "tax_unit_taxable_social_security", + "adjusted_gross_income", + "income_tax", + "employee_social_security_tax", + "employer_social_security_tax", + "self_employment_social_security_tax", + "ssi_unearned_income", + "ssi", + "snap_unearned_income", + "snap", + "medicare_irmaa_tax_year_magi", + "medicare_irmaa_magi_two_years_prior", + "medicare_irmaa_magi", + "medicare_part_b_premium", +) +OUTPUTS = (*PREEXISTING_SOCIAL_SECURITY, *DOWNSTREAM) +NAMES = ("retiree", "disabled") +AMOUNTS = ((24_000, 0, 1_800, 600), (0, 6_000, 1_200, 300)) + + +class HouseholdMicrosimulation(Microsimulation): + default_dataset = None + + +@pytest.fixture(scope="module") +def original_input_system(): + """Restore main's input contracts, including effective per-capita growth.""" + original = system.clone() + for public_name in SOCIAL_SECURITY_COMPONENTS: + reported = original.variables[f"{public_name}_reported"] + input_class = type( + public_name, + (Variable,), + { + "value_type": reported.value_type, + "entity": reported.entity, + "definition_period": YEAR, + "unit": reported.unit, + "label": public_name, + "uprating": reported.uprating, + }, + ) + original.replace_variable(input_class) + restored = original.variables[public_name] + assert restored.is_input_variable() + assert restored.uprating == reported.uprating + assert restored.uprating.endswith("social_security_per_capita") + return original + + +def situation(*, source="public", years=(2024, 2026), include_benefits=True): + people = {} + for index, name in enumerate(NAMES): + person = { + "age": {year: (70 if index == 0 else 45) + year - 2024 for year in years}, + "employment_income": {2024: 15_000 if index == 0 else 0}, + "self_employment_income": {2024: 3_000 if index == 0 else 0}, + "taxable_interest_income": {2024: 50_000 if index == 0 else 10}, + "is_disabled": {2024: index == 1}, + "takes_up_ssi_if_eligible": {2024: True}, + } + if include_benefits: + for public_name, amount in zip(SOCIAL_SECURITY_COMPONENTS, AMOUNTS[index]): + name_to_store = ( + f"{public_name}_reported" if source == "reported" else public_name + ) + person[name_to_store] = { + year: amount * (1 if year == 2024 else 1.07) for year in years + } + people[name] = person + result = {"people": people} + for plural in ("households", "tax_units", "spm_units", "families", "marital_units"): + result[plural] = {name: {"members": [name]} for name in NAMES} + for household in result["households"].values(): + household["state_code"] = {year: "TX" for year in years} + for spm_unit in result["spm_units"].values(): + spm_unit["takes_up_snap_if_eligible"] = {2024: True} + return result + + +def reference_simulation(original_input_system, wrapper=Simulation, **kwargs): + # Both wrappers retain all their existing income/geography preparation. + # Only this PR's compatibility migration is absent from the reference. + with patch( + "policyengine_us.system.migrate_social_security_inputs", return_value=set() + ): + return wrapper(tax_benefit_system=original_input_system, **kwargs) + + +def assert_equal_outputs(current, original, requests, variables=OUTPUTS): + for period in requests: + for variable in variables: + np.testing.assert_array_equal( + np.asarray(current.calculate(variable, period)), + np.asarray(original.calculate(variable, period)), + err_msg=f"Current-law difference for {variable}, {period}", + ) + + +@pytest.mark.parametrize( + "requests", + [ + ("2024-12", 2026, "2025-03", 2024, 2025, "2026-08"), + ("2026-08", 2025, 2024, "2025-03", 2026, "2024-12"), + ], + ids=("month-first-forward", "month-first-reverse"), +) +def test_household_current_law_matches_original_through_years_and_months( + original_input_system, requests +): + inputs = situation() + original = reference_simulation(original_input_system, situation=deepcopy(inputs)) + current = Simulation(situation=inputs, tax_benefit_system=system) + assert_equal_outputs(current, original, requests) + + # The supplied 2024 source must be read again for 2026 IRMAA, even after + # calculating 2026 income and a different source amount for that year. + np.testing.assert_array_equal( + current.calculate("medicare_irmaa_magi_two_years_prior", 2026), + current.calculate("medicare_irmaa_tax_year_magi", 2024), + ) + assert not np.array_equal( + current.calculate("medicare_irmaa_tax_year_magi", 2026), + current.calculate("medicare_irmaa_tax_year_magi", 2024), + ) + + +def dataset_frame(): + """The dataset loader itself sets the stored public columns as inputs.""" + data = {"person_id": [1, 2]} + for entity in ("household", "tax_unit", "spm_unit", "family", "marital_unit"): + data[f"person_{entity}_id"] = [1, 2] + data[f"{entity}_id"] = [1, 2] + for year in (2024, 2026): + data[f"age__{year}"] = [70 + year - 2024, 45 + year - 2024] + data[f"state_fips__{year}"] = [48, 48] + for component_index, public_name in enumerate(SOCIAL_SECURITY_COMPONENTS): + data[f"{public_name}__{year}"] = [ + amounts[component_index] * (1 if year == 2024 else 1.07) + for amounts in AMOUNTS + ] + data["employment_income__2024"] = [15_000, 0] + data["self_employment_income__2024"] = [3_000, 0] + data["taxable_interest_income__2024"] = [50_000, 10] + data["is_disabled__2024"] = [False, True] + data["takes_up_ssi_if_eligible__2024"] = [True, True] + data["takes_up_snap_if_eligible__2024"] = [True, True] + return pd.DataFrame(data) + + +def test_dataset_loader_current_law_matches_original(original_input_system): + frame = dataset_frame() + original = reference_simulation( + original_input_system, Microsimulation, dataset=frame.copy(deep=True) + ) + current = Microsimulation(dataset=frame, tax_benefit_system=system) + assert_equal_outputs(current, original, ("2025-06", 2026, 2024, 2025)) + + +@pytest.mark.parametrize("ingress", ["core_builder", "set_input"]) +def test_core_builder_and_dataset_style_source_writes_match_original( + original_input_system, ingress +): + if ingress == "core_builder": + original_situation = situation() + current_situation = situation(source="reported") + else: + original_situation = situation(include_benefits=False) + current_situation = deepcopy(original_situation) + original = SimulationBuilder().build_from_dict( + original_input_system, original_situation + ) + current = SimulationBuilder().build_from_dict(system, current_situation) + if ingress == "set_input": + for component_index, public_name in enumerate(SOCIAL_SECURITY_COMPONENTS): + for year in (2024, 2026): + values = [ + amounts[component_index] * (1 if year == 2024 else 1.07) + for amounts in AMOUNTS + ] + original.set_input(public_name, year, values) + current.set_input(f"{public_name}_reported", year, values) + assert_equal_outputs(current, original, ("2026-11", 2025, 2024, 2026)) + + +def test_clones_and_neutral_reform_branches_match_original(original_input_system): + original = reference_simulation(original_input_system, situation=situation()) + current = Simulation(situation=situation(), tax_benefit_system=system) + assert_equal_outputs(current, original, (2024, 2026)) + + # Warm-cache clones and ordinary branches should preserve the source + # uprating path as well as any final public override applied later. + for current_copy, original_copy in ( + ( + current.clone(clone_tax_benefit_system=False), + original.clone(clone_tax_benefit_system=False), + ), + (current.get_branch("copy"), original.get_branch("copy")), + ): + current_copy.set_input("social_security_retirement", 2026, [0, 123]) + original_copy.set_input("social_security_retirement", 2026, [0, 123]) + drop_inherited_values(current_copy) + drop_inherited_values(original_copy) + assert_equal_outputs(current_copy, original_copy, (2026, "2025-07")) + + # Register an actual no-change parameter reform after outputs were + # cached. This exercises reform replay and invalidation, while copying + # policy only for these two branches and preserving the read-only models. + neutral = Reform.from_dict( + {"gov.contrib.ssa.benefit_levers.across_the_board.rate": {"2024": 0}} + ) + original_reform = original.get_branch("neutral_reform", clone_system=True) + current_reform = current.get_branch("neutral_reform", clone_system=True) + original_reform.apply_reform(neutral) + current_reform.apply_reform(neutral) + assert_equal_outputs(current_reform, original_reform, (2025, 2026)) + assert_equal_outputs(current, original, (2024, 2026)) + + +def test_legacy_extraction_reload_matches_original(original_input_system): + inputs = situation(years=(2024,)) + original = reference_simulation( + original_input_system, HouseholdMicrosimulation, situation=deepcopy(inputs) + ) + current = HouseholdMicrosimulation(situation=inputs, tax_benefit_system=system) + original_payload = original.extract_person(0) + current_payload = current.extract_person(0) + original_reload = reference_simulation( + original_input_system, situation=original_payload + ) + current_reload = Simulation(situation=current_payload, tax_benefit_system=system) + assert_equal_outputs(current_reload, original_reload, ("2024-10", 2024, 2025)) + for public_name, amount in zip(SOCIAL_SECURITY_COMPONENTS, AMOUNTS[0]): + assert current_reload.calculate(public_name, 2024)[0] == amount + + +def test_complete_source_and_final_override_reload_matches_original( + original_input_system, +): + original = reference_simulation(original_input_system, situation=situation()) + current = Simulation(situation=situation(), tax_benefit_system=system) + # An unsupplied, calculated year must remain a calculation after export. + assert_equal_outputs(current, original, (2025,)) + for simulation in (original, current): + simulation.set_input("social_security_retirement", 2027, [0, 123]) + simulation.set_input("social_security", 2027, [9_999, 0]) + original_payload = export_source_inputs(original) + current_payload = export_source_inputs(current) + for public_name in SOCIAL_SECURITY_COMPONENTS: + reported = current_payload["sources"]["people"]["retiree"][ + f"{public_name}_reported" + ] + assert set(reported) == {"2024", "2026"} + with patch( + "policyengine_us.system.migrate_social_security_inputs", return_value=set() + ): + original_reload = reload_source_inputs( + original_payload, tax_benefit_system=original_input_system + ) + current_reload = reload_source_inputs(current_payload, tax_benefit_system=system) + requests = ("2027-08", 2026, 2024, 2025, 2027) + assert_equal_outputs(current_reload, original_reload, requests) + assert_equal_outputs(current_reload, current, requests) + + +def test_later_reported_write_updates_a_fresh_branch_without_changing_parent( + original_input_system, +): + current = Simulation(situation=situation(), tax_benefit_system=system) + original = reference_simulation(original_input_system, situation=situation()) + assert_equal_outputs(current, original, (2024,)) + parent_keys = set(current._user_input_keys) + current_branch = current.get_branch("updated_reported") + original_branch = original.get_branch("updated_original") + current_branch.set_input("social_security_retirement_reported", 2024, [30000, 2400]) + original_branch.set_input("social_security_retirement", 2024, [30000, 2400]) + drop_inherited_values(current_branch) + drop_inherited_values(original_branch) + assert_equal_outputs(current_branch, original_branch, (2024, 2025)) + assert current._user_input_keys == parent_keys + np.testing.assert_array_equal( + current.calculate("social_security_retirement", 2024), [24000, 0] + ) + reloaded = reload_source_inputs( + export_source_inputs(current_branch), tax_benefit_system=system + ) + assert_equal_outputs(reloaded, current_branch, (2024, 2025)) + + +@pytest.mark.parametrize("wrapper", [Simulation, HouseholdMicrosimulation]) +def test_constructor_created_baseline_keeps_migrated_uprating( + original_input_system, + wrapper, +): + # Core creates its automatic baseline before wrapper migration. A source + # present only in 2024 must still uprate into 2025 on that baseline. + neutral = Reform.from_dict( + {"gov.contrib.ssa.benefit_levers.across_the_board.rate": {"2024": 0}} + ) + inputs = situation(years=(2024,)) + original = reference_simulation( + original_input_system, wrapper, situation=deepcopy(inputs) + ) + current = wrapper(situation=inputs, tax_benefit_system=system, reform=neutral) + assert current.baseline is not None + assert_equal_outputs(current.baseline, original, (2024, "2025-06", 2025)) + assert_equal_outputs(current, original, (2024, 2025)) + for public in SOCIAL_SECURITY_COMPONENTS: + assert not current.baseline._get_set_input_periods(public) + assert current.baseline._get_set_input_periods(f"{public}_reported") diff --git a/policyengine_us/tests/core/test_social_security_inputs.py b/policyengine_us/tests/core/test_social_security_inputs.py new file mode 100644 index 00000000000..ed0b5d03b2a --- /dev/null +++ b/policyengine_us/tests/core/test_social_security_inputs.py @@ -0,0 +1,305 @@ +"""Input compatibility, branch isolation and complete source round trips. + +These cases need Python because YAML bypasses PE-US constructors and cannot +exercise cloning, stored-input bookkeeping or source export. They reuse the +read-only country model and use one person, without loading population data. +""" + +import json + +import numpy as np +import pandas as pd +import pytest +from policyengine_core.simulations import SimulationBuilder + +from policyengine_us import Microsimulation, Simulation +from policyengine_us.data.economic_assumptions import MICRODATA_UPRATING_OVERRIDES +from policyengine_us.system import system +from policyengine_us.tools.period_branch import drop_inherited_values +from policyengine_us.tools.social_security_inputs import ( + SOCIAL_SECURITY_COMPONENTS, + export_source_inputs, + reload_source_inputs, +) + + +class HouseholdMicrosimulation(Microsimulation): + """Exercise the actual wrapper with a household instead of its default data.""" + + default_dataset = None + + +WRAPPERS = (Simulation, HouseholdMicrosimulation) +FACTOR = "social_security_across_the_board_factor" + + +def situation(inputs): + return { + "people": {"person": {"age": {2024: 70, 2025: 71, 2026: 72}, **inputs}}, + "households": {"household": {"members": ["person"]}}, + "tax_units": {"tax_unit": {"members": ["person"]}}, + "spm_units": {"spm_unit": {"members": ["person"]}}, + "families": {"family": {"members": ["person"]}}, + "marital_units": {"marital_unit": {"members": ["person"]}}, + } + + +def make_simulation(wrapper=Simulation, inputs=None): + return wrapper(situation=situation(inputs or {}), tax_benefit_system=system) + + +def amount(simulation, name, period): + return np.asarray(simulation.calculate(name, period))[0] + + +def input_periods(simulation, name): + return {str(period) for period in simulation._get_set_input_periods(name)} + + +@pytest.mark.parametrize("wrapper", WRAPPERS) +def test_constructor_precedence_is_per_component_and_period(wrapper): + inputs = {} + for index, public in enumerate(SOCIAL_SECURITY_COMPONENTS, 1): + inputs[public] = {2024: index * 1_000, 2026: index * 3_000} + inputs[f"{public}_reported"] = {2025: index * 1_100, 2026: index * 1_200} + simulation = make_simulation(wrapper, inputs) + for year in (2024, 2025, 2026): + simulation.set_input(FACTOR, year, [1.1]) + for index, public in enumerate(SOCIAL_SECURITY_COMPONENTS, 1): + reported = f"{public}_reported" + assert input_periods(simulation, public) == {"2026"} + assert input_periods(simulation, reported) == {"2024", "2025", "2026"} + assert amount(simulation, public, 2024) == pytest.approx(index * 1_100) + assert amount(simulation, public, 2025) == pytest.approx(index * 1_210) + assert amount(simulation, reported, 2026) == index * 1_200 + assert amount(simulation, public, 2026) == index * 3_000 + assert public in simulation.input_variables + assert reported in simulation.input_variables + + +@pytest.mark.parametrize("wrapper", WRAPPERS) +def test_constructor_zero_and_aggregate_are_explicit_inputs(wrapper): + inputs = { + SOCIAL_SECURITY_COMPONENTS[0]: {2024: 0}, + f"{SOCIAL_SECURITY_COMPONENTS[1]}_reported": {2024: 0}, + SOCIAL_SECURITY_COMPONENTS[2]: {2024: 0}, + f"{SOCIAL_SECURITY_COMPONENTS[2]}_reported": {2024: 500}, + "social_security": {2024: 123}, + } + simulation = make_simulation(wrapper, inputs) + simulation.set_input(FACTOR, 2024, [1.5]) + for public in SOCIAL_SECURITY_COMPONENTS: + assert amount(simulation, public, 2024) == 0 + assert amount(simulation, "social_security", 2024) == 123 + assert SOCIAL_SECURITY_COMPONENTS[0] not in simulation.input_variables + assert f"{SOCIAL_SECURITY_COMPONENTS[0]}_reported" in simulation.input_variables + + +@pytest.mark.parametrize("wrapper", WRAPPERS) +def test_later_public_override_survives_clones_and_cache_clear(wrapper): + public = SOCIAL_SECURITY_COMPONENTS[0] + simulation = make_simulation(wrapper, {public: {2024: 12_000}}) + assert amount(simulation, public, 2024) == 12_000 + simulation.set_input(public, 2024, [0]) + original_keys = set(simulation._user_input_keys) + for copy in ( + simulation.clone(clone_tax_benefit_system=False), + simulation.get_branch("override"), + ): + drop_inherited_values(copy) + assert amount(copy, public, 2024) == 0 + copy.set_input(public, 2024, [42]) + drop_inherited_values(copy) + assert amount(copy, public, 2024) == 42 + assert simulation._user_input_keys == original_keys + assert amount(simulation, public, 2024) == 0 + assert amount(simulation, f"{public}_reported", 2024) == 12_000 + + +def test_core_builder_public_inputs_are_final_overrides(): + retirement, disability = SOCIAL_SECURITY_COMPONENTS[:2] + simulation = SimulationBuilder().build_from_dict( + system, + situation( + { + f"{retirement}_reported": {2024: 12_000}, + disability: {2024: 777}, + FACTOR: {2024: 1.1}, + } + ), + ) + assert amount(simulation, retirement, 2024) == pytest.approx(13_200) + assert amount(simulation, disability, 2024) == 777 + assert input_periods(simulation, disability) == {"2024"} + simulation.set_input(retirement, 2024, [0]) + assert amount(simulation, retirement, 2024) == 0 + + +def test_dataset_constructor_migrates_stored_original_columns(): + public = SOCIAL_SECURITY_COMPONENTS[0] + frame = pd.DataFrame( + { + "person_id": [1], + **{ + f"person_{entity}_id": [1] + for entity in ( + "household", + "tax_unit", + "spm_unit", + "family", + "marital_unit", + ) + }, + **{ + f"{entity}_id": [1] + for entity in ( + "household", + "tax_unit", + "spm_unit", + "family", + "marital_unit", + ) + }, + f"{public}__2024": [12_000], + f"{public}__2026": [14_000], + f"{public}_reported__2026": [13_000], + } + ) + simulation = Microsimulation(dataset=frame, tax_benefit_system=system) + simulation.set_input(FACTOR, 2024, [1.1]) + assert amount(simulation, public, 2024) == pytest.approx(13_200) + assert amount(simulation, public, 2026) == 14_000 + assert input_periods(simulation, public) == {"2026"} + assert input_periods(simulation, f"{public}_reported") == {"2024", "2026"} + # Loader-time public writes migrate; later writes remain final outputs. + simulation.set_input(public, 2024, [0]) + assert amount(simulation, public, 2024) == 0 + assert amount(simulation, f"{public}_reported", 2024) == 12_000 + reloaded = reload_source_inputs( + export_source_inputs(simulation), tax_benefit_system=system + ) + assert amount(reloaded, public, 2024) == 0 + assert amount(reloaded, f"{public}_reported", 2024) == 12_000 + + +def test_complete_snapshot_preserves_histories_and_separate_final_overrides(): + public = SOCIAL_SECURITY_COMPONENTS[0] + simulation = make_simulation( + inputs={ + f"{public}_reported": {2024: 12_000, 2026: 14_000}, + public: {2026: 6_500}, + f"{public}_source": {2024: "reported", 2026: "reported"}, + "ss_aime_input": {2024: 1_000, 2026: 2_000}, + "employment_income": {2024: 10_000, 2026: 20_000}, + } + ) + # A calculated intervening year must never become a reported source. + amount(simulation, public, 2025) + simulation.set_input(public, 2027, [0]) + payload = json.loads(json.dumps(export_source_inputs(simulation))) + person = payload["sources"]["people"]["person"] + assert set(person[f"{public}_reported"]) == {"2024", "2026"} + assert public not in person + assert set(payload["overrides"][public]) == {"2026", "2027"} + assert person["ss_aime_input"] == {"2024": 1_000, "2026": 2_000} + reloaded = reload_source_inputs(payload, tax_benefit_system=system) + for name in (public, f"{public}_reported", "ss_aime_input", "employment_income"): + for year in (2024, 2025, 2026, 2027): + assert amount(reloaded, name, year) == pytest.approx( + amount(simulation, name, year) + ) + assert input_periods(reloaded, public) == {"2026", "2027"} + assert export_source_inputs(reloaded)["overrides"][public]["2027"] == [0] + + +def test_legacy_extraction_uses_reported_sources_and_months_keep_flow_conversion(): + public = SOCIAL_SECURITY_COMPONENTS[0] + simulation = make_simulation(HouseholdMicrosimulation, {public: {2024: 12_000}}) + extracted = simulation.extract_person(0) + person = next(iter(extracted["people"].values())) + assert f"{public}_reported" in person + assert public not in person + reloaded = Simulation(situation=extracted, tax_benefit_system=system) + for month_first in (False, True): + copy = reloaded.clone(clone_tax_benefit_system=False) + if month_first: + assert amount(copy, public, "2024-07") == 1_000 + assert amount(copy, public, 2024) == 12_000 + assert amount(copy, public, "2024-07") == 1_000 + + +def test_reported_twins_are_pure_and_original_dataset_names_keep_their_uprater(): + for public in SOCIAL_SECURITY_COMPONENTS: + original = system.get_variable(public) + reported = system.get_variable(f"{public}_reported") + assert not original.is_input_variable() + assert original.uprating is None + assert reported.is_input_variable() + assert not reported.formulas and not reported.adds and not reported.subtracts + assert reported.uprating.endswith("social_security_per_capita") + assert ( + MICRODATA_UPRATING_OVERRIDES[public] + == "calibration.gov.irs.soi.social_security" + ) + + +@pytest.mark.parametrize( + "source", ["supplied_aime", "supplied_pia", "supplied_ratio_survivor", "history"] +) +def test_reserved_sources_fail_explicitly_instead_of_using_reported_fallback(source): + public = SOCIAL_SECURITY_COMPONENTS[0] + simulation = make_simulation( + inputs={f"{public}_source": {2024: source}, f"{public}_reported": {2024: 0}} + ) + with pytest.raises(NotImplementedError, match="reserved"): + simulation.calculate(public, 2024) + + +def test_oasi_only_cut_rejects_unidentified_positive_dependent_stream(): + from policyengine_core.reforms import Reform + + reform = Reform.from_dict( + { + "gov.contrib.ssa.benefit_levers.payable.share": {"2024": 0.8}, + "gov.contrib.ssa.benefit_levers.payable.oasi_only": {"2024": True}, + } + ) + simulation = Simulation( + situation=situation({"social_security_dependents_reported": {2024: 500}}), + reform=reform, + tax_benefit_system=system, + ) + with pytest.raises(ValueError, match="dependent benefit fund"): + simulation.calculate("social_security_dependents", 2024) + + +def test_parameter_reform_registration_order_preserves_stable_public_formula(): + from policyengine_core.reforms import Reform + + across = Reform.from_dict( + { + "gov.contrib.ssa.benefit_levers.across_the_board.start_year": { + "2024": 2024 + }, + "gov.contrib.ssa.benefit_levers.across_the_board.rate": {"2024": 0.1}, + } + ) + payable = Reform.from_dict( + {"gov.contrib.ssa.benefit_levers.payable.share": {"2024": 0.8}} + ) + first = Simulation( + situation=situation({"social_security_retirement_reported": {2024: 12000}}), + reform=(across, payable), + tax_benefit_system=system, + ) + second = Simulation( + situation=situation({"social_security_retirement_reported": {2024: 12000}}), + reform=(payable, across), + tax_benefit_system=system, + ) + assert amount(first, "social_security_retirement", 2024) == pytest.approx(10560) + assert amount(second, "social_security_retirement", 2024) == pytest.approx(10560) + assert ( + first.tax_benefit_system.variables["social_security_retirement"].formulas + == second.tax_benefit_system.variables["social_security_retirement"].formulas + ) diff --git a/policyengine_us/tests/core/test_social_security_lever_properties.py b/policyengine_us/tests/core/test_social_security_lever_properties.py new file mode 100644 index 00000000000..f2488eb7766 --- /dev/null +++ b/policyengine_us/tests/core/test_social_security_lever_properties.py @@ -0,0 +1,399 @@ +"""Pure COLA algebra and rounding properties required for benefit levers. + +Python is needed for generated ledgers, ordering, and Decimal boundaries. +These tests construct no simulations; public-variable policy examples belong +in YAML, with ingress and migration checked by the separate integration tests. +""" + +from math import prod + +import pytest +from hypothesis import given, settings, strategies as st + +from policyengine_us.tools.social_security_benefit_levers import ( + joint_cola_factor, + projected_cola_rates, +) + + +@st.composite +def cola_ledgers(draw): + first_year = draw(st.integers(min_value=1984, max_value=2050)) + count = draw(st.integers(min_value=1, max_value=12)) + rates = draw( + st.lists( + st.integers(min_value=0, max_value=1500), + min_size=count, + max_size=count, + ) + ) + changes = draw( + st.lists( + st.integers(min_value=-2000, max_value=2000), + min_size=count, + max_size=count, + ) + ) + withholding = draw(st.lists(st.booleans(), min_size=count, max_size=count)) + events = [ + { + "payment_year": first_year + offset, + "payment_month": 1, + "rate": rate / 10000, + "rounding_rule": 2, + } + for offset, rate in enumerate(rates) + ] + first_payment = ( + first_year + draw(st.integers(min_value=-2, max_value=count + 1)), + draw(st.integers(min_value=1, max_value=12)), + ) + through_year = first_year + draw(st.integers(min_value=-2, max_value=count + 1)) + pp_changes = { + first_year + offset: change / 100000 for offset, change in enumerate(changes) + } + withheld = {first_year + offset: value for offset, value in enumerate(withholding)} + return events, first_payment, through_year, pp_changes, withheld + + +@settings(max_examples=100, deadline=None) +@given(cola_ledgers()) +def test_joint_factor_is_nonnegative(ledger): + assert joint_cola_factor(*ledger) >= 0 + + +@settings(max_examples=100, deadline=None) +@given( + cola_ledgers(), + st.integers(min_value=-2000, max_value=2000), + st.integers(min_value=-2000, max_value=2000), +) +def test_joint_factor_is_monotone_in_cola_change(ledger, delta_a, delta_b): + events, first_payment, through_year, _, withheld = ledger + lower, upper = sorted((delta_a / 100000, delta_b / 100000)) + lower_changes = {event["payment_year"]: lower for event in events} + upper_changes = {event["payment_year"]: upper for event in events} + lower_factor = joint_cola_factor( + events, first_payment, through_year, lower_changes, withheld + ) + upper_factor = joint_cola_factor( + events, first_payment, through_year, upper_changes, withheld + ) + assert lower_factor <= upper_factor + 1e-12 + + +@settings(max_examples=100, deadline=None) +@given(cola_ledgers()) +def test_unexposed_benefit_has_unit_factor(ledger): + events, _, through_year, pp_changes, withheld = ledger + after_last_event = (events[-1]["payment_year"] + 1, 1) + assert ( + joint_cola_factor(events, after_last_event, through_year, pp_changes, withheld) + == 1 + ) + + +@settings(max_examples=100, deadline=None) +@given(cola_ledgers(), st.integers(min_value=0, max_value=12)) +def test_factor_is_one_before_change_start(ledger, start_offset): + events, first_payment, _, pp_changes, _ = ledger + start_year = events[0]["payment_year"] + start_offset + future_changes = { + year: change for year, change in pp_changes.items() if year >= start_year + } + assert ( + joint_cola_factor(events, first_payment, start_year - 1, future_changes, {}) + == 1 + ) + + +@settings(max_examples=100, deadline=None) +@given(cola_ledgers()) +def test_event_factors_multiply_and_event_order_commutes(ledger): + events, first_payment, through_year, pp_changes, withheld = ledger + whole = joint_cola_factor(*ledger) + separate = prod( + joint_cola_factor([event], first_payment, through_year, pp_changes, withheld) + for event in events + ) + reversed_factor = joint_cola_factor( + list(reversed(events)), + first_payment, + through_year, + pp_changes, + withheld, + ) + assert whole == pytest.approx(separate, rel=1e-12) + assert whole == pytest.approx(reversed_factor, rel=1e-12) + + +@settings(max_examples=100, deadline=None) +@given( + cola_ledgers(), + st.integers(min_value=0, max_value=1000000), + st.integers(min_value=-1000, max_value=2000), + st.integers(min_value=0, max_value=1000), +) +def test_independent_factors_commute_and_reported_zero_stays_zero( + ledger, reported, across_rate_points, payable_points +): + cola = joint_cola_factor(*ledger) + across = 1 + across_rate_points / 1000 + payable = payable_points / 1000 + assert across >= 0 + assert payable >= 0 + combined = reported * prod((cola, across, payable)) + assert combined >= 0 + assert combined == pytest.approx((reported * payable) * across * cola, rel=1e-12) + assert 0 * prod((cola, across, payable)) == 0 + + +def test_projection_uses_high_water_after_decline_and_rounded_zero(): + rates = projected_cola_rates( + { + 2024: 100, + 2025: 99, + 2026: 100.04, + 2027: 100.05, + 2028: 101.0505, + } + ) + assert rates[2025] == 0 + assert rates[2026] == 0 + # 0.05 percent ties round upward to 0.1 percent, rather than to even. + assert rates[2027] == 0.001 + # Once a positive COLA occurs, its Q3 index becomes the comparison base. + assert rates[2028] == 0.01 + + +def test_projection_rounds_statutory_rate_instead_of_using_raw_index_ratio(): + rates = projected_cola_rates({2025: 308.729, 2026: 317.265}) + assert rates[2026] == 0.028 + + +@pytest.mark.parametrize("delta", [0.0015, -0.0030]) +def test_presets_keep_half_tenth_precision_after_baseline_rounding(delta): + event = { + "payment_year": 2027, + "payment_month": 1, + "rate": 0.028, + "rounding_rule": 2, + } + factor = joint_cola_factor([event], (2020, 1), 2027, {2027: delta}, {}) + assert factor == pytest.approx((1.028 + delta) / 1.028) + + +def test_withheld_year_uses_joint_numerator_and_unchanged_denominator(): + event = { + "payment_year": 2027, + "payment_month": 1, + "rate": 0.028, + "rounding_rule": 2, + } + factor = joint_cola_factor([event], (2020, 1), 2027, {2027: 0.0015}, {2027: True}) + assert factor == pytest.approx(1 / 1.028) + + +def test_changed_cola_is_floored_at_zero(): + event = { + "payment_year": 2027, + "payment_month": 1, + "rate": 0.001, + "rounding_rule": 2, + } + factor = joint_cola_factor([event], (2020, 1), 2027, {2027: -0.0030}, {}) + assert factor == pytest.approx(1 / 1.001) + + +def test_historical_july_event_respects_first_payment_month(): + event = { + "payment_year": 1981, + "payment_month": 7, + "rate": 0.112, + "rounding_rule": 1, + } + exposed = joint_cola_factor([event], (1981, 7), 1981, {1981: 0.0015}, {}) + assert exposed == pytest.approx(1.1135 / 1.112) + assert joint_cola_factor([event], (1981, 8), 1981, {1981: 0.0015}, {}) == 1 + + +def test_unsupported_rounding_is_rejected_only_for_exposed_history(): + event = { + "payment_year": 2027, + "payment_month": 1, + "rate": 0.028, + "rounding_rule": 999, + } + with pytest.raises(ValueError, match="rounding"): + joint_cola_factor([event], (2020, 1), 2027, {}, {}) + assert joint_cola_factor([event], (2028, 1), 2027, {}, {}) == 1 + + +@pytest.mark.parametrize( + "variable,field,value", + [ + ("social_security_across_the_board_factor", "rate", -1.01), + ("social_security_across_the_board_factor", "rate", float("inf")), + ("social_security_across_the_board_factor", "rate", float("nan")), + ("social_security_payable_factor", "share", -0.01), + ("social_security_payable_factor", "share", 1.01), + ("social_security_payable_factor", "share", float("inf")), + ("social_security_payable_factor", "share", float("nan")), + ], +) +def test_invalid_policy_factors_are_rejected(variable, field, value): + # YAML cannot assert exceptions. Call the authored formula with a minimal + # parameter view; no additional country model or simulation is built. + from types import SimpleNamespace + from policyengine_core.periods import period + from policyengine_us.system import system + + parameter = SimpleNamespace(**{field: value, "start_year": 2024}) + leaf = "across_the_board" if field == "rate" else "payable" + tree = SimpleNamespace( + gov=SimpleNamespace( + contrib=SimpleNamespace( + ssa=SimpleNamespace(benefit_levers=SimpleNamespace(**{leaf: parameter})) + ) + ) + ) + formula = system.variables[variable].get_formula("2024-01-01") + with pytest.raises(ValueError): + formula(None, period("2024"), lambda _: tree) + + +def exposure_parameters(**overrides): + from policyengine_core.parameters import ParameterNode + from policyengine_us.system import system + + values = dict( + start_year=9999, + effective_year=0, + effective_month=12, + end_year=0, + pp_change=0, + alternative_rate=-1, + withholding_enabled=False, + ) + values.update(overrides) + cola = { + name: {"values": value if isinstance(value, dict) else {"0000-01-01": value}} + for name, value in values.items() + } + params = ParameterNode( + data={ + "gov": { + "contrib": {"ssa": {"benefit_levers": {"cola": cola}}}, + "ssa": { + "uprating": { + "values": { + "2026-01-01": 317.265, + "2027-01-01": 328.746, + "2028-01-01": 332.241, + } + } + }, + } + } + ) + params.gov.ssa.add_child( + "cola_history", system.parameters.gov.ssa.cola_history.clone() + ) + return params + + +class ExposurePerson: + ids = ["worker"] + + def __init__(self, age, **inputs): + self.age = age + self.inputs = inputs + + def __call__(self, name, year): + import numpy as np + from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, + ) + + if name.endswith("_source"): + return SocialSecurityBenefitSource.encode(np.array(["reported"])) + if name == "age": + return np.array([self.age]) + if name.endswith("_worker_id"): + return np.array([""]) + if name.endswith("_month"): + return np.array([1]) + value = self.inputs.get(name, 0) + if isinstance(value, dict): + value = value.get(int(str(year)[:4]), False) + return np.array([value]) + + +def test_historical_june_change_exposes_retirement_eligibility_year(): + from policyengine_core.periods import period + from policyengine_us.tools.social_security_benefit_levers import ( + reported_joint_cola_factor, + ) + + params = exposure_parameters( + effective_year=1981, effective_month=6, pp_change=0.0015 + ) + factor = reported_joint_cola_factor( + ExposurePerson(62), period("1981"), params, "retirement" + ) + assert factor[0] == pytest.approx(1.1135 / 1.112) + + +def test_prior_withheld_cola_remains_after_withholding_is_disabled(): + from policyengine_core.periods import period + from policyengine_us.tools.social_security_benefit_levers import ( + reported_joint_cola_factor, + ) + + params = exposure_parameters( + withholding_enabled={ + "0000-01-01": False, + "2027-01-01": True, + "2028-01-01": False, + } + ) + person = ExposurePerson( + 70, + social_security_retirement_first_cola_year=2027, + social_security_retirement_cola_withheld={2027: True}, + ) + factor = reported_joint_cola_factor(person, period("2028"), params, "retirement") + assert factor[0] == pytest.approx(1 / 1.036) + + +def test_mixed_calculated_rows_bypass_population_exposure_validation(): + import numpy as np + from policyengine_core.periods import period + from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, + reported_joint_cola_factor, + ) + + class MixedPerson: + ids = ["reported", "calculated"] + + def __call__(self, name, year): + if name.endswith("_source"): + return SocialSecurityBenefitSource.encode( + np.array(["reported", "supplied_pia"]) + ) + if name.endswith("_worker_id"): + return np.array(["", "irrelevant-missing-record"]) + if name.endswith("_month"): + return np.array([1, 13]) + if name.endswith("_first_cola_year"): + return np.array([2027, 2027]) + if name == "age": + return np.array([70, 70]) + return np.array([0, 0]) + + params = exposure_parameters(start_year=2027, pp_change=0.0015) + factors = reported_joint_cola_factor( + MixedPerson(), period("2027"), params, "retirement" + ) + np.testing.assert_allclose(factors, [1.0375 / 1.036, 1]) diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_dependents_joint_cola_factor.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_dependents_joint_cola_factor.yaml new file mode 100644 index 00000000000..ec795529bdb --- /dev/null +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_dependents_joint_cola_factor.yaml @@ -0,0 +1,57 @@ +- name: Positive child benefits use the explicit pre-start exposure proxy + period: 2027 + input: + age: 8 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + social_security_dependents_reported: 6000 + social_security_survivors_reported: 6000 + social_security_disability_reported: 6000 + output: + social_security_dependents_joint_cola_factor: 1.0014478764478765 + social_security_survivors_joint_cola_factor: 1.0014478764478765 + social_security_disability_joint_cola_factor: 1.0014478764478765 + +- name: Newly entitled child sensitivity is a separately supplied scenario + period: 2027 + input: + age: 8 + social_security_dependents_entitlement_year: 2027 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + social_security_dependents_reported: 6000 + output: + social_security_dependents_joint_cola_factor: 1 + +- name: Reported zero remains zero under combined levers + period: 2027 + input: + age: 8 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + gov.contrib.ssa.benefit_levers.across_the_board.start_year: 2027 + gov.contrib.ssa.benefit_levers.across_the_board.rate: 0.1 + gov.contrib.ssa.benefit_levers.payable.share: 0.8 + social_security_dependents_reported: 0 + output: + social_security_dependents: 0 + +- name: Supplied worker link takes precedence over a young recipient proxy + period: 2027 + input: + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + people: + worker: + age: 62 + social_security_worker_computation_year: 2027 + child: + age: 8 + social_security_dependents_worker_id: worker + social_security_dependents_reported: 6000 + output: + people: + child: + social_security_dependents_joint_cola_factor: 1 + worker: + social_security_dependents_joint_cola_factor: 1.0014478764478765 diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement.yaml new file mode 100644 index 00000000000..d659a6ea1ef --- /dev/null +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement.yaml @@ -0,0 +1,127 @@ +- name: Reported components retain current-law annual amounts + period: 2027 + input: + social_security_retirement_reported: 12000 + social_security_disability_reported: 6000 + social_security_survivors_reported: 3000 + social_security_dependents_reported: 1000 + output: + social_security_retirement: 12000 + social_security_disability: 6000 + social_security_survivors: 3000 + social_security_dependents: 1000 + social_security: 22000 + +- name: Independent across-the-board and payable factors multiply once + period: 2027 + input: + gov.contrib.ssa.benefit_levers.across_the_board.start_year: 2027 + gov.contrib.ssa.benefit_levers.across_the_board.rate: 0.1 + gov.contrib.ssa.benefit_levers.payable.share: 0.8 + social_security_retirement_reported: 12000 + social_security_disability_reported: 6000 + social_security_survivors_reported: 3000 + social_security_dependents_reported: 1000 + output: + social_security_retirement: 10560 + social_security_disability: 5280 + social_security_survivors: 2640 + social_security_dependents: 880 + social_security: 19360 + +- name: Across-the-board change begins in its payment year + period: 2026 + input: + gov.contrib.ssa.benefit_levers.across_the_board.start_year: 2027 + gov.contrib.ssa.benefit_levers.across_the_board.rate: 0.1 + social_security_retirement_reported: 12000 + output: + social_security_retirement: 12000 + +- name: OASI-only payable share preserves disability benefits + period: 2027 + input: + gov.contrib.ssa.benefit_levers.payable.share: 0.8 + gov.contrib.ssa.benefit_levers.payable.oasi_only: true + social_security_dependents_fund: oasi + social_security_retirement_reported: 12000 + social_security_disability_reported: 6000 + social_security_survivors_reported: 3000 + social_security_dependents_reported: 1000 + output: + social_security_retirement: 9600 + social_security_disability: 6000 + social_security_survivors: 2400 + social_security_dependents: 800 + +- name: A full across-the-board cut produces zero + period: 2027 + input: + gov.contrib.ssa.benefit_levers.across_the_board.start_year: 2027 + gov.contrib.ssa.benefit_levers.across_the_board.rate: -1 + social_security_retirement_reported: 12000 + output: + social_security_retirement: 0 + +- name: Explicit public and aggregate inputs remain final overrides in core YAML + period: 2027 + input: + gov.contrib.ssa.benefit_levers.across_the_board.start_year: 2027 + gov.contrib.ssa.benefit_levers.across_the_board.rate: 0.1 + social_security_retirement_reported: 12000 + social_security_retirement: 0 + social_security: 777 + output: + social_security_retirement: 0 + social_security: 777 + +- name: A COLA change composes with across-the-board and payable changes + period: 2027 + absolute_error_margin: 0.01 + input: + age: 70 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + gov.contrib.ssa.benefit_levers.across_the_board.start_year: 2027 + gov.contrib.ssa.benefit_levers.across_the_board.rate: 0.1 + gov.contrib.ssa.benefit_levers.payable.share: 0.8 + social_security_retirement_reported: 12000 + output: + social_security_retirement: 10575.289575289577 + +- name: OASI-only payable scenario preserves a dependent DI benefit + period: 2027 + input: + gov.contrib.ssa.benefit_levers.payable.share: 0.8 + gov.contrib.ssa.benefit_levers.payable.oasi_only: true + social_security_dependents_fund: di + social_security_dependents_reported: 6000 + output: + social_security_dependents: 6000 + +- name: Calculated-base extension point bypasses population PIA and COLA factors + period: 2027 + input: + age: 70 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + gov.contrib.ssa.benefit_levers.across_the_board.start_year: 2027 + gov.contrib.ssa.benefit_levers.across_the_board.rate: 0.1 + gov.contrib.ssa.benefit_levers.payable.share: 0.8 + social_security_retirement_source: supplied_pia + # The route itself is reserved. A supplied selected base stands in for + # the future service's amount already containing its reformed PIA/COLAs. + social_security_retirement_base: 10000 + social_security_retirement_pia_ratio_factor: 1.2 + social_security_retirement_joint_cola_factor: 1.3 + output: + social_security_retirement: 8800 + +- name: Selecting an OASI-only basis before depletion remains neutral + period: 2027 + input: + gov.contrib.ssa.benefit_levers.payable.share: 1 + gov.contrib.ssa.benefit_levers.payable.oasi_only: true + social_security_dependents_reported: 6000 + output: + social_security_dependents: 6000 diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_joint_cola_factor.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_joint_cola_factor.yaml new file mode 100644 index 00000000000..c6588bce46a --- /dev/null +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_joint_cola_factor.yaml @@ -0,0 +1,118 @@ +- name: CPI-E precision survives the statutory baseline rounding + period: 2027 + absolute_error_margin: 0.01 + input: + age: 70 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + social_security_retirement_reported: 12000 + output: + # January 2027 baseline is 3.6%; 12000 * 1.0375 / 1.036. + social_security_retirement_joint_cola_factor: 1.0014478764478765 + social_security_retirement: 12017.374517374517 + +- name: Newly eligible retirement has no same-year COLA exposure + period: 2027 + input: + age: 62 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + social_security_retirement_reported: 12000 + output: + social_security_retirement_joint_cola_factor: 1 + social_security_retirement: 12000 + +- name: Static age cross-sections infer eligibility separately in each year + period: 2035 + input: + age: 70 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.end_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + social_security_retirement_reported: 12000 + output: + # This row's inferred eligibility is 2027, first COLA January 2028. + social_security_retirement_joint_cola_factor: 1 + social_security_retirement: 12000 + +- name: Retirement COLAs accrue before a later claiming age + period: 2027 + input: + age: 67 + ss_claiming_age: 67 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + social_security_retirement_reported: 12000 + output: + social_security_retirement_joint_cola_factor: 1.0014478764478765 + +- name: Supplied first-COLA date overrides an older-age proxy + period: 2027 + input: + age: 75 + social_security_retirement_first_cola_year: 2028 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + social_security_retirement_reported: 12000 + output: + social_security_retirement_joint_cola_factor: 1 + social_security_retirement: 12000 + +- name: Withheld COLA uses one numerator even when CPI-E is active + period: 2027 + absolute_error_margin: 0.01 + input: + age: 70 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + gov.contrib.ssa.benefit_levers.cola.withholding_enabled: true + social_security_retirement_first_cola_year: 2027 + social_security_retirement_cola_withheld: true + social_security_retirement_reported: 12000 + output: + social_security_retirement_joint_cola_factor: 0.9652509652509652 + social_security_retirement: 11583.011583011583 + +- name: Chained CPI change retains its requested precision + period: 2027 + input: + age: 70 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: -0.0030 + social_security_retirement_reported: 12000 + output: + social_security_retirement_joint_cola_factor: 0.997104247104247 + +- name: December effective date first changes the following January check + period: 2028 + input: + age: 70 + gov.contrib.ssa.benefit_levers.cola.effective_year: 2028 + gov.contrib.ssa.benefit_levers.cola.effective_month: 12 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + social_security_retirement_reported: 12000 + output: + social_security_retirement_joint_cola_factor: 1 + +- name: First affected January after December effective date + period: 2029 + input: + age: 70 + gov.contrib.ssa.benefit_levers.cola.effective_year: 2028 + gov.contrib.ssa.benefit_levers.cola.effective_month: 12 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + social_security_retirement_reported: 12000 + output: + # Statutory baseline 339.431 / 332.241 - 1 rounds to 2.2%. + social_security_retirement_joint_cola_factor: 1.0014677103718199 + +- name: A retirement claim date does not erase eligibility COLAs + period: 2027 + input: + age: 70 + social_security_retirement_entitlement_year: 2027 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + social_security_retirement_reported: 12000 + output: + social_security_retirement_joint_cola_factor: 1.0014478764478765 diff --git a/policyengine_us/tools/social_security_benefit_levers.py b/policyengine_us/tools/social_security_benefit_levers.py new file mode 100644 index 00000000000..fb86fa14647 --- /dev/null +++ b/policyengine_us/tools/social_security_benefit_levers.py @@ -0,0 +1,257 @@ +"""Payment-event algebra and the reported Social Security benefit route. + +This module never derives COLAs from benefit-income uprating. Historical events +are statutory rates; projections use the system's January-keyed Q3 schedule. +""" + +from decimal import Decimal, ROUND_HALF_UP + +import numpy as np +from policyengine_core.model_api import Enum + +from policyengine_us.tools.ssa_cola_history import validate_cola_rounding_rule + + +class SocialSecurityBenefitSource(Enum): + reported = "Reported benefits" + supplied_aime = "Supplied eligibility-level AIME" + supplied_pia = "Supplied unreformed eligibility-level PIA" + supplied_ratio_survivor = "Supplied-ratio survivor calculation" + history = "Annual earnings history" + + +class SocialSecurityFund(Enum): + unspecified = "Underlying trust fund not supplied" + oasi = "Old-Age and Survivors Insurance" + di = "Disability Insurance" + + +def projected_cola_rates(indices): + """Return statutory rates after the first (anchor) January index year. + + A zero rounded increase does not reset the last effective COLA's Q3 base. + """ + years = sorted(indices) + if not years: + return {} + high_water = Decimal(str(indices[years[0]])) + if not high_water.is_finite() or high_water <= 0: + raise ValueError("SSA Q3 index levels must be finite and positive") + rates = {} + for year in years[1:]: + index = Decimal(str(indices[year])) + if not index.is_finite() or index <= 0: + raise ValueError("SSA Q3 index levels must be finite and positive") + rate = max(Decimal(0), index / high_water - 1).quantize( + Decimal("0.001"), rounding=ROUND_HALF_UP + ) + rates[year] = float(rate) + if rate > 0: + high_water = index + return rates + + +def joint_cola_factor(events, first_payment, through_year, pp_changes, withheld): + """One numerator per exposed event, including the withholding decision. + + Annual reported amounts use the payment-calendar-year convention. First + payment month distinguishes the historical July events from January checks. + This ratio is a population approximation, not a reconstructed rounded PIA. + """ + factor = 1.0 + for event in events: + year, month = event["payment_year"], event["payment_month"] + if (year, month) < first_payment or year > through_year: + continue + validate_cola_rounding_rule(event["rounding_rule"]) + rate = float(event["rate"]) + delta = float(pp_changes.get(year, 0)) + factor *= cola_event_multiplier(rate, delta, withheld.get(year, False)) + return float(factor) + + +def cola_event_multiplier(rate, delta, withheld): + """Shared scalar/vector algebra for an exposed payment event.""" + if not np.isfinite(rate) or rate < 0 or not np.isfinite(delta): + raise ValueError("COLA rates must be nonnegative and changes finite") + numerator = np.where(withheld, 1, 1 + max(0, rate + delta)) + return numerator / (1 + rate) + + +def cola_events(parameters, through_year): + """Export the exact historical and projected payment-event schedule used.""" + history = parameters.gov.ssa.cola_history + events = [] + for entry in reversed(history.rate.values_list): + instant = entry.instant_str + p = history(instant) + effective_year = int(instant[:4]) + payment_year = effective_year + (p.effective_month == 12) + events.append( + dict( + effective_date=instant, + payment_year=int(payment_year), + payment_month=int(p.first_payment_month), + rate=float(entry.value), + rounding_rule=int(p.pia_rounding), + source="historical", + ) + ) + last_actual = max(e["payment_year"] for e in events) + if through_year > last_actual: + indices = { + year: parameters(f"{year}-01-01").gov.ssa.uprating + for year in range(last_actual, through_year + 1) + } + for year, rate in projected_cola_rates(indices).items(): + events.append( + dict( + effective_date=f"{year - 1}-12-01", + payment_year=year, + payment_month=1, + rate=rate, + rounding_rule=2, + source="gov.ssa.uprating", + ) + ) + return sorted( + (e for e in events if e["payment_year"] <= through_year), + key=lambda e: (e["payment_year"], e["payment_month"]), + ) + + +def selected_reported_benefit(person, period, component): + prefix = f"social_security_{component}" + source = person(f"{prefix}_source", period) + if np.any(source != SocialSecurityBenefitSource.reported): + raise NotImplementedError( + "Calculated Social Security source modes are reserved; " + "this implementation supports the reported route only" + ) + return person(f"{prefix}_reported", period) + + +def reported_joint_cola_factor(person, period, parameters, component): + prefix = f"social_security_{component}" + reported = ( + person(f"{prefix}_source", period) == SocialSecurityBenefitSource.reported + ) + if not np.any(reported): + return 1 + year = period.start.year + p = parameters(period).gov.contrib.ssa.benefit_levers.cola + if not 1 <= p.effective_month <= 12: + raise ValueError("COLA effective month must lie in [1, 12]") + raw_policy = parameters.gov.contrib.ssa.benefit_levers.cola + # An earlier withheld or adjusted COLA remains in the payment level after + # a policy is switched off. Inspect the parameter history, not just today. + active = any( + entry.instant_str <= f"{year}-12-31" + and (entry.value >= 0 if name == "alternative_rate" else bool(entry.value)) + for name in ("withholding_enabled", "pp_change", "alternative_rate") + for entry in getattr(raw_policy, name).values_list + ) + if not active: + return 1 + events = cola_events(parameters, year) + first_year = person(f"{prefix}_first_cola_year", period) + first_month = person(f"{prefix}_first_cola_month", period) + entitlement_month = person(f"{prefix}_entitlement_month", period) + if np.any(reported & ((first_month < 1) | (first_month > 12))) or np.any( + reported & ((entitlement_month < 1) | (entitlement_month > 12)) + ): + raise ValueError("Social Security input months must lie in [1, 12]") + worker_year = person(f"{prefix}_worker_eligibility_year", period) + worker_id = person(f"{prefix}_worker_id", period) + shared_worker_year = person("social_security_worker_computation_year", period) + # Links follow an explicitly supplied record basis, never a recipient-age + # retirement assumption on an auxiliary/disability record. + ids = person.ids + known_workers = {str(identifier): i for i, identifier in enumerate(ids)} + linked = worker_year.copy() + for i, identifier in enumerate(worker_id): + if not reported[i]: + continue + if str(identifier) in known_workers and identifier != "": + worker_index = known_workers[str(identifier)] + record_year = worker_year[worker_index] or shared_worker_year[worker_index] + if linked[i] and record_year and linked[i] != record_year: + raise ValueError( + "Conflicting supplied and linked worker computation years" + ) + linked[i] = linked[i] or record_year + if linked[i] == 0 and first_year[i] <= 0: + raise ValueError("Linked worker requires an explicit eligibility year") + elif identifier != "": + raise ValueError("Social Security worker_id does not identify a person") + entitlement = person(f"{prefix}_entitlement_year", period) + age = person("age", period) + if component == "retirement": + eligibility = year - age + 62 + proxy = eligibility + (eligibility >= 1983) + elif component == "survivors": + eligibility = year - age + 60 + proxy = np.where(age >= 60, eligibility + (eligibility >= 1983), 1975) + else: + proxy = np.full(len(age), 1975) + # Modern December COLAs first reach January checks in the next year. + # For supplied historical entitlement dates, June-effective events are + # included only once the recipient was entitled by that effective month. + historical_first = np.where( + (entitlement > 0) & (entitlement <= 1982), + entitlement + (entitlement_month > 6), + entitlement + 1, + ) + # A retirement entitlement date can be a late claim. It must not discard + # the COLAs accrued since eligibility at 62. + recipient_origin = ( + proxy + if component == "retirement" + else np.where(entitlement > 0, historical_first, proxy) + ) + inferred = np.where(linked > 0, linked + (linked >= 1983), recipient_origin) + first_year = np.where(first_year > 0, first_year, inferred) + result = np.ones(len(age)) + for event in events: + payment_year, month = event["payment_year"], event["payment_month"] + exposed = (first_year < payment_year) | ( + (first_year == payment_year) & (first_month <= month) + ) + exposed &= reported + if not np.any(exposed): + continue + validate_cola_rounding_rule(event["rounding_rule"]) + policy = parameters(f"{payment_year}-01-01").gov.contrib.ssa.benefit_levers.cola + effective = event["effective_date"] + applies = ( + ( + effective + >= f"{int(policy.effective_year):04d}-{int(policy.effective_month):02d}-01" + ) + if policy.effective_year + else payment_year >= policy.start_year + ) + applies &= not policy.end_year or payment_year <= policy.end_year + baseline = event["rate"] + alternative = policy.alternative_rate + if ( + not np.isfinite(policy.pp_change) + or not np.isfinite(alternative) + or (alternative < 0 and alternative != -1) + ): + raise ValueError( + "COLA change must be finite; alternative rate must be nonnegative or -1" + ) + rate = ( + (alternative if alternative >= 0 else baseline) + policy.pp_change + if applies + else baseline + ) + withheld = ( + person(f"{prefix}_cola_withheld", payment_year) + if policy.withholding_enabled + else False + ) + multiplier = cola_event_multiplier(baseline, rate - baseline, withheld) + result *= np.where(exposed, multiplier, 1) + return result diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py new file mode 100644 index 00000000000..59fca399937 --- /dev/null +++ b/policyengine_us/tools/social_security_inputs.py @@ -0,0 +1,152 @@ +"""Compatibility ingress and complete, source-safe simulation snapshots. + +Constructor migration is intentionally separate from ``set_input``: a later +write to a public benefit is an output override, including an explicit zero. +""" + +import json + +from policyengine_core.simulations.simulation import NpEncoder + +from policyengine_us.tools.period_branch import delete_input_values + + +SOCIAL_SECURITY_COMPONENTS = ( + "social_security_retirement", + "social_security_disability", + "social_security_survivors", + "social_security_dependents", +) + + +def migrate_social_security_inputs(simulation, *, include_baseline=True): + """Move legacy component inputs, period by period, into pure sources. + + A reported input in the same visible period makes the public input a + deliberate final override. Inputs in other periods do not affect that + choice. Recorded-input helpers exclude formula caches and honor visible + branch precedence. + """ + visible_branches = simulation._get_visible_branch_names() + migrated = set() + for public_name in SOCIAL_SECURITY_COMPONENTS: + reported_name = f"{public_name}_reported" + reported_periods = set(simulation._get_set_input_periods(reported_name)) + holder = simulation.get_holder(public_name) + for period in simulation._get_set_input_periods(public_name): + if period in reported_periods: + continue + value = simulation._get_recorded_input_array(public_name, period) + if value is None: + continue + simulation.set_input(reported_name, period, value.copy()) + for branch_name in visible_branches: + delete_input_values(holder, period, branch_name) + migrated.add(public_name) + + # Core's legacy extraction reads this runtime list, which was assembled + # before migration. Retain public names only where actual overrides remain. + runtime_inputs = [] + for name in simulation.input_variables: + if name not in migrated or simulation._get_set_input_periods(name): + runtime_inputs.append(name) + if name in migrated: + runtime_inputs.append(f"{name}_reported") + for public_name in SOCIAL_SECURITY_COMPONENTS: + reported_name = f"{public_name}_reported" + if simulation._get_set_input_periods(reported_name): + runtime_inputs.append(reported_name) + simulation.input_variables = list(dict.fromkeys(runtime_inputs)) + # Core creates this clone during super().__init__, before either PE-US + # wrapper reaches migration. Its copied public sources otherwise override + # formulas in the supplied year and disappear from the uprating route in + # later years. Each branch needs its own source/input-key migration. + baseline = getattr(simulation, "baseline", None) + if include_baseline and baseline is not None and baseline is not simulation: + migrate_social_security_inputs(baseline, include_baseline=False) + return migrated + + +def export_source_inputs(simulation): + """Export all supplied periods, separating pure sources and final overrides. + + The JSON-compatible payload preserves entity memberships and every recorded + pure-input period, including selectors, dates, worker links and earnings + histories. Supplied formula-variable inputs have their own ``overrides`` + payload; calculated results never become sources. Only the current branch + and its visible ancestors are exported, using their normal precedence. + + This is the complete-history alternative to Core's legacy one-period + ``extract_person`` API. Export the desired small simulation or household; + this helper deliberately does not materialize a population dataset. + """ + names = { + key: [str(name) for name in population.ids] + for key, population in simulation.populations.items() + } + sources = {} + for key, population in simulation.populations.items(): + entity = population.entity + sources[entity.plural] = {name: {} for name in names[key]} + if not entity.is_person: + # All PE-US group entities use the unbounded member role. + for person_index, group_index in enumerate(population.members_entity_id): + group = sources[entity.plural][names[key][int(group_index)]] + group.setdefault("members", []).append(names["person"][person_index]) + + overrides = {} + for name, variable in simulation.tax_benefit_system.variables.items(): + periods = simulation._get_set_input_periods(name) + if not periods: + continue + for period in periods: + values = simulation._get_recorded_input_array(name, period) + if values is None: + continue + if hasattr(values, "decode_to_str"): + values = values.decode_to_str() + if variable.is_input_variable(): + entity = variable.entity + for index, entity_name in enumerate(names[entity.key]): + source = sources[entity.plural][entity_name] + source.setdefault(name, {})[str(period)] = values[index] + else: + overrides.setdefault(name, {})[str(period)] = values + + return json.loads( + json.dumps( + { + "version": 1, + "sources": sources, + "overrides": overrides, + "default_input_period": str(simulation.default_input_period), + "default_calculation_period": str( + simulation.default_calculation_period + ), + }, + cls=NpEncoder, + ) + ) + + +def reload_source_inputs(payload, **simulation_kwargs): + """Reload a complete source snapshot, then restore its final overrides. + + Pass the desired ``tax_benefit_system`` or ``reform`` through keyword + arguments. Loading overrides after construction preserves a public-only + final override instead of treating it as a legacy reported source. + """ + from policyengine_us import Simulation + + if payload.get("version") != 1: + raise ValueError("Unsupported source-input snapshot version") + for period_name in ("default_input_period", "default_calculation_period"): + simulation_kwargs.setdefault(period_name, payload[period_name]) + simulation = Simulation(situation=payload["sources"], **simulation_kwargs) + for name, periods in payload["overrides"].items(): + for period, values in periods.items(): + simulation.set_input(name, period, values) + simulation.input_variables = list( + dict.fromkeys([*simulation.input_variables, *payload["overrides"]]) + ) + return simulation diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_across_the_board_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_across_the_board_factor.py new file mode 100644 index 00000000000..d5c65cc9b14 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_across_the_board_factor.py @@ -0,0 +1,16 @@ +from policyengine_us.model_api import * + + +class social_security_across_the_board_factor(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Social Security across-the-board factor" + + def formula(person, period, parameters): + p = parameters(period).gov.contrib.ssa.benefit_levers.across_the_board + if not np.isfinite(p.rate) or p.rate < -1: + raise ValueError( + "Social Security across-the-board rate must be finite and at least -1" + ) + return 1 + p.rate if period.start.year >= p.start_year else 1 diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents.py index eee9679d417..17a31f93045 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_dependents.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents.py @@ -1,4 +1,7 @@ from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, +) class social_security_dependents(Variable): @@ -7,4 +10,18 @@ class social_security_dependents(Variable): definition_period = YEAR label = "Social Security dependents benefits" unit = USD - uprating = "calibration.gov.irs.soi.social_security" + documentation = "Selected source with applicable benefit levers; an explicit public input is a final override." + + def formula(person, period, parameters): + return ( + person("social_security_dependents_base", period) + * where( + person("social_security_dependents_source", period) + == SocialSecurityBenefitSource.reported, + person("social_security_dependents_joint_cola_factor", period) + * person("social_security_dependents_pia_ratio_factor", period), + 1, + ) + * person("social_security_across_the_board_factor", period) + * person("social_security_dependents_payable_factor", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_base.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_base.py new file mode 100644 index 00000000000..3f471a2be4b --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_base.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + selected_reported_benefit, +) + + +class social_security_dependents_base(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = "Selected Social Security dependents benefit base" + + def formula(person, period, parameters): + return selected_reported_benefit(person, period, "dependents") diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_cola_withheld.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_cola_withheld.py new file mode 100644 index 00000000000..99af5c2342a --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_cola_withheld.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_cola_withheld(Variable): + value_type = bool + default_value = False + entity = Person + definition_period = YEAR + label = "Social Security dependents cola withheld" + documentation = "Whether this payment year COLA is withheld; inactive unless withholding_enabled is set." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_entitlement_month.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_entitlement_month.py new file mode 100644 index 00000000000..75e9442bb94 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_entitlement_month.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_entitlement_month(Variable): + value_type = int + default_value = 1 + entity = Person + definition_period = YEAR + label = "Social Security dependents entitlement month" + documentation = "Recipient entitlement month (1 through 12)." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_entitlement_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_entitlement_year.py new file mode 100644 index 00000000000..78dddd3dc7b --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_entitlement_year.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_entitlement_year(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security dependents entitlement year" + documentation = "Recipient entitlement year; zero means unavailable." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_first_cola_month.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_first_cola_month.py new file mode 100644 index 00000000000..d0d3ade903c --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_first_cola_month.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_first_cola_month(Variable): + value_type = int + default_value = 1 + entity = Person + definition_period = YEAR + label = "Social Security dependents first cola month" + documentation = "First exposed COLA payment month (1 through 12)." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_first_cola_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_first_cola_year.py new file mode 100644 index 00000000000..af7655efad9 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_first_cola_year.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_first_cola_year(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security dependents first cola year" + documentation = "First exposed COLA payment year; zero requests the documented population proxy." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_fund.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_fund.py new file mode 100644 index 00000000000..7fc1a3a5dc0 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_fund.py @@ -0,0 +1,12 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import SocialSecurityFund + + +class social_security_dependents_fund(Variable): + value_type = Enum + possible_values = SocialSecurityFund + default_value = SocialSecurityFund.unspecified + entity = Person + definition_period = YEAR + label = "Underlying Social Security dependent benefit trust fund" + documentation = "OASI-only cuts require a supplied OASI or DI basis for a positive dependent benefit. Mixed streams require a supplied blended scenario rather than an invented allocation." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_joint_cola_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_joint_cola_factor.py new file mode 100644 index 00000000000..64539ed253b --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_joint_cola_factor.py @@ -0,0 +1,17 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + reported_joint_cola_factor, +) + + +class social_security_dependents_joint_cola_factor(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Joint Social Security dependents COLA factor" + documentation = ( + "One payment-event numerator combines adjusted COLAs and withholding." + ) + + def formula(person, period, parameters): + return reported_joint_cola_factor(person, period, parameters, "dependents") diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_payable_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_payable_factor.py new file mode 100644 index 00000000000..b1beac8c468 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_payable_factor.py @@ -0,0 +1,22 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import SocialSecurityFund + + +class social_security_dependents_payable_factor(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Social Security dependents payable factor" + + def formula(person, period, parameters): + p = parameters(period).gov.contrib.ssa.benefit_levers.payable + factor = person("social_security_payable_factor", period) + if not p.oasi_only or np.all(factor == 1): + return factor + fund = person("social_security_dependents_fund", period) + base = person("social_security_dependents_base", period) + if np.any((base > 0) & (fund == SocialSecurityFund.unspecified)): + raise ValueError( + "OASI-only payable cuts require the dependent benefit fund" + ) + return where(fund == SocialSecurityFund.oasi, factor, 1) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_factor.py new file mode 100644 index 00000000000..99543e62b56 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_factor.py @@ -0,0 +1,14 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_pia_ratio_factor(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Social Security dependents PIA ratio factor" + documentation = ( + "Neutral extension point for a subsequent reported-benefit PIA reform." + ) + + def formula(person, period, parameters): + return 1 diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_reported.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_reported.py new file mode 100644 index 00000000000..a545607a6d6 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_reported.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_reported(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Social Security dependents benefits, reported" + unit = USD + uprating = "calibration.gov.irs.soi.social_security" diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_source.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_source.py new file mode 100644 index 00000000000..a5dd86b34d7 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_source.py @@ -0,0 +1,13 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, +) + + +class social_security_dependents_source(Variable): + value_type = Enum + possible_values = SocialSecurityBenefitSource + default_value = SocialSecurityBenefitSource.reported + entity = Person + definition_period = YEAR + label = "Social Security dependents benefit source" diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_eligibility_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_eligibility_year.py new file mode 100644 index 00000000000..cca4a2541fe --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_eligibility_year.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_worker_eligibility_year(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security dependents worker eligibility year" + documentation = "Explicit underlying-worker eligibility or computation year; zero means unavailable." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_id.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_id.py new file mode 100644 index 00000000000..33e13f74f47 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_id.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_worker_id(Variable): + value_type = str + default_value = "" + entity = Person + definition_period = YEAR + label = "Social Security dependents worker id" + documentation = "Underlying-worker person identifier; requires an explicit computation year on that record." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability.py index 00c3a5f70de..540f39139a5 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_disability.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability.py @@ -1,10 +1,27 @@ from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, +) class social_security_disability(Variable): value_type = float entity = Person definition_period = YEAR - label = "Social Security disability benefits (SSDI)" + label = "Social Security disability benefits" unit = USD - uprating = "calibration.gov.irs.soi.social_security" + documentation = "Selected source with applicable benefit levers; an explicit public input is a final override." + + def formula(person, period, parameters): + return ( + person("social_security_disability_base", period) + * where( + person("social_security_disability_source", period) + == SocialSecurityBenefitSource.reported, + person("social_security_disability_joint_cola_factor", period) + * person("social_security_disability_pia_ratio_factor", period), + 1, + ) + * person("social_security_across_the_board_factor", period) + * person("social_security_disability_payable_factor", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_base.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_base.py new file mode 100644 index 00000000000..394433a8919 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_base.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + selected_reported_benefit, +) + + +class social_security_disability_base(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = "Selected Social Security disability benefit base" + + def formula(person, period, parameters): + return selected_reported_benefit(person, period, "disability") diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_cola_withheld.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_cola_withheld.py new file mode 100644 index 00000000000..96208e14665 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_cola_withheld.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_disability_cola_withheld(Variable): + value_type = bool + default_value = False + entity = Person + definition_period = YEAR + label = "Social Security disability cola withheld" + documentation = "Whether this payment year COLA is withheld; inactive unless withholding_enabled is set." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_entitlement_month.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_entitlement_month.py new file mode 100644 index 00000000000..d8c6d9e15c3 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_entitlement_month.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_disability_entitlement_month(Variable): + value_type = int + default_value = 1 + entity = Person + definition_period = YEAR + label = "Social Security disability entitlement month" + documentation = "Recipient entitlement month (1 through 12)." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_entitlement_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_entitlement_year.py new file mode 100644 index 00000000000..fdf96f65c8f --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_entitlement_year.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_disability_entitlement_year(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security disability entitlement year" + documentation = "Recipient entitlement year; zero means unavailable." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_first_cola_month.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_first_cola_month.py new file mode 100644 index 00000000000..350a0f78a1a --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_first_cola_month.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_disability_first_cola_month(Variable): + value_type = int + default_value = 1 + entity = Person + definition_period = YEAR + label = "Social Security disability first cola month" + documentation = "First exposed COLA payment month (1 through 12)." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_first_cola_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_first_cola_year.py new file mode 100644 index 00000000000..fed2bf8b691 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_first_cola_year.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_disability_first_cola_year(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security disability first cola year" + documentation = "First exposed COLA payment year; zero requests the documented population proxy." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_joint_cola_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_joint_cola_factor.py new file mode 100644 index 00000000000..e2a73b44adf --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_joint_cola_factor.py @@ -0,0 +1,17 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + reported_joint_cola_factor, +) + + +class social_security_disability_joint_cola_factor(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Joint Social Security disability COLA factor" + documentation = ( + "One payment-event numerator combines adjusted COLAs and withholding." + ) + + def formula(person, period, parameters): + return reported_joint_cola_factor(person, period, parameters, "disability") diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_payable_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_payable_factor.py new file mode 100644 index 00000000000..efbfaafb98e --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_payable_factor.py @@ -0,0 +1,13 @@ +from policyengine_us.model_api import * + + +class social_security_disability_payable_factor(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Social Security disability payable factor" + + def formula(person, period, parameters): + p = parameters(period).gov.contrib.ssa.benefit_levers.payable + factor = person("social_security_payable_factor", period) + return 1 if p.oasi_only and "disability" == "disability" else factor diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_factor.py new file mode 100644 index 00000000000..3d641e7cf60 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_factor.py @@ -0,0 +1,14 @@ +from policyengine_us.model_api import * + + +class social_security_disability_pia_ratio_factor(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Social Security disability PIA ratio factor" + documentation = ( + "Neutral extension point for a subsequent reported-benefit PIA reform." + ) + + def formula(person, period, parameters): + return 1 diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_reported.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_reported.py new file mode 100644 index 00000000000..504f198f8c7 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_reported.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_disability_reported(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Social Security disability benefits (SSDI)" + unit = USD + uprating = "calibration.gov.irs.soi.social_security" diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_source.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_source.py new file mode 100644 index 00000000000..2aecb8307e4 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_source.py @@ -0,0 +1,13 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, +) + + +class social_security_disability_source(Variable): + value_type = Enum + possible_values = SocialSecurityBenefitSource + default_value = SocialSecurityBenefitSource.reported + entity = Person + definition_period = YEAR + label = "Social Security disability benefit source" diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_worker_eligibility_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_worker_eligibility_year.py new file mode 100644 index 00000000000..d94933c3ffe --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_worker_eligibility_year.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_disability_worker_eligibility_year(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security disability worker eligibility year" + documentation = "Explicit underlying-worker eligibility or computation year; zero means unavailable." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_worker_id.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_worker_id.py new file mode 100644 index 00000000000..253fdeeab23 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_worker_id.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_disability_worker_id(Variable): + value_type = str + default_value = "" + entity = Person + definition_period = YEAR + label = "Social Security disability worker id" + documentation = "Underlying-worker person identifier; requires an explicit computation year on that record." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_payable_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_payable_factor.py new file mode 100644 index 00000000000..9f653f8f065 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_payable_factor.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * + + +class social_security_payable_factor(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Social Security supplied payable-share factor" + documentation = "Externally supplied annual share, including a supplier-blended depletion year; no depletion projection is inferred." + + def formula(person, period, parameters): + p = parameters(period).gov.contrib.ssa.benefit_levers.payable + if not np.isfinite(p.share) or not 0 <= p.share <= 1: + raise ValueError("Social Security payable share must lie in [0, 1]") + return p.share diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement.py index 076caf2a9bc..9918ec46f7b 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_retirement.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement.py @@ -1,4 +1,7 @@ from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, +) class social_security_retirement(Variable): @@ -7,4 +10,18 @@ class social_security_retirement(Variable): definition_period = YEAR label = "Social Security retirement benefits" unit = USD - uprating = "calibration.gov.irs.soi.social_security" + documentation = "Selected source with applicable benefit levers; an explicit public input is a final override." + + def formula(person, period, parameters): + return ( + person("social_security_retirement_base", period) + * where( + person("social_security_retirement_source", period) + == SocialSecurityBenefitSource.reported, + person("social_security_retirement_joint_cola_factor", period) + * person("social_security_retirement_pia_ratio_factor", period), + 1, + ) + * person("social_security_across_the_board_factor", period) + * person("social_security_retirement_payable_factor", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_base.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_base.py new file mode 100644 index 00000000000..3f55abcb9f3 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_base.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + selected_reported_benefit, +) + + +class social_security_retirement_base(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = "Selected Social Security retirement benefit base" + + def formula(person, period, parameters): + return selected_reported_benefit(person, period, "retirement") diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_cola_withheld.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_cola_withheld.py new file mode 100644 index 00000000000..c448d954234 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_cola_withheld.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_cola_withheld(Variable): + value_type = bool + default_value = False + entity = Person + definition_period = YEAR + label = "Social Security retirement cola withheld" + documentation = "Whether this payment year COLA is withheld; inactive unless withholding_enabled is set." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_entitlement_month.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_entitlement_month.py new file mode 100644 index 00000000000..97be0c068f9 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_entitlement_month.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_entitlement_month(Variable): + value_type = int + default_value = 1 + entity = Person + definition_period = YEAR + label = "Social Security retirement entitlement month" + documentation = "Recipient entitlement month (1 through 12)." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_entitlement_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_entitlement_year.py new file mode 100644 index 00000000000..db84e7221ed --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_entitlement_year.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_entitlement_year(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security retirement entitlement year" + documentation = "Recipient entitlement year; zero means unavailable." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_first_cola_month.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_first_cola_month.py new file mode 100644 index 00000000000..381401705f8 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_first_cola_month.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_first_cola_month(Variable): + value_type = int + default_value = 1 + entity = Person + definition_period = YEAR + label = "Social Security retirement first cola month" + documentation = "First exposed COLA payment month (1 through 12)." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_first_cola_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_first_cola_year.py new file mode 100644 index 00000000000..bba4b81da2c --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_first_cola_year.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_first_cola_year(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security retirement first cola year" + documentation = "First exposed COLA payment year; zero requests the documented population proxy." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_joint_cola_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_joint_cola_factor.py new file mode 100644 index 00000000000..55214555c25 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_joint_cola_factor.py @@ -0,0 +1,17 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + reported_joint_cola_factor, +) + + +class social_security_retirement_joint_cola_factor(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Joint Social Security retirement COLA factor" + documentation = ( + "One payment-event numerator combines adjusted COLAs and withholding." + ) + + def formula(person, period, parameters): + return reported_joint_cola_factor(person, period, parameters, "retirement") diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_payable_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_payable_factor.py new file mode 100644 index 00000000000..f88697a58ba --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_payable_factor.py @@ -0,0 +1,13 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_payable_factor(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Social Security retirement payable factor" + + def formula(person, period, parameters): + p = parameters(period).gov.contrib.ssa.benefit_levers.payable + factor = person("social_security_payable_factor", period) + return 1 if p.oasi_only and "retirement" == "disability" else factor diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_factor.py new file mode 100644 index 00000000000..fd3609292ac --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_factor.py @@ -0,0 +1,14 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_pia_ratio_factor(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Social Security retirement PIA ratio factor" + documentation = ( + "Neutral extension point for a subsequent reported-benefit PIA reform." + ) + + def formula(person, period, parameters): + return 1 diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_reported.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_reported.py new file mode 100644 index 00000000000..5a89d1cbfc9 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_reported.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_reported(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Social Security retirement benefits, reported" + unit = USD + uprating = "calibration.gov.irs.soi.social_security" diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_source.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_source.py new file mode 100644 index 00000000000..446bbe603db --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_source.py @@ -0,0 +1,13 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, +) + + +class social_security_retirement_source(Variable): + value_type = Enum + possible_values = SocialSecurityBenefitSource + default_value = SocialSecurityBenefitSource.reported + entity = Person + definition_period = YEAR + label = "Social Security retirement benefit source" diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_worker_eligibility_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_worker_eligibility_year.py new file mode 100644 index 00000000000..a3e471d22db --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_worker_eligibility_year.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_worker_eligibility_year(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security retirement worker eligibility year" + documentation = "Explicit underlying-worker eligibility or computation year; zero means unavailable." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_worker_id.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_worker_id.py new file mode 100644 index 00000000000..f39aa49ffb4 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_worker_id.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_worker_id(Variable): + value_type = str + default_value = "" + entity = Person + definition_period = YEAR + label = "Social Security retirement worker id" + documentation = "Underlying-worker person identifier; requires an explicit computation year on that record." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors.py index 5120158e549..78d1a8505ff 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_survivors.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors.py @@ -1,4 +1,7 @@ from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, +) class social_security_survivors(Variable): @@ -7,4 +10,18 @@ class social_security_survivors(Variable): definition_period = YEAR label = "Social Security survivors benefits" unit = USD - uprating = "calibration.gov.irs.soi.social_security" + documentation = "Selected source with applicable benefit levers; an explicit public input is a final override." + + def formula(person, period, parameters): + return ( + person("social_security_survivors_base", period) + * where( + person("social_security_survivors_source", period) + == SocialSecurityBenefitSource.reported, + person("social_security_survivors_joint_cola_factor", period) + * person("social_security_survivors_pia_ratio_factor", period), + 1, + ) + * person("social_security_across_the_board_factor", period) + * person("social_security_survivors_payable_factor", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_base.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_base.py new file mode 100644 index 00000000000..2f2ba569c74 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_base.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + selected_reported_benefit, +) + + +class social_security_survivors_base(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = "Selected Social Security survivors benefit base" + + def formula(person, period, parameters): + return selected_reported_benefit(person, period, "survivors") diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_cola_withheld.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_cola_withheld.py new file mode 100644 index 00000000000..0e151f3d6da --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_cola_withheld.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_cola_withheld(Variable): + value_type = bool + default_value = False + entity = Person + definition_period = YEAR + label = "Social Security survivors cola withheld" + documentation = "Whether this payment year COLA is withheld; inactive unless withholding_enabled is set." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_entitlement_month.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_entitlement_month.py new file mode 100644 index 00000000000..82faf0f0c5b --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_entitlement_month.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_entitlement_month(Variable): + value_type = int + default_value = 1 + entity = Person + definition_period = YEAR + label = "Social Security survivors entitlement month" + documentation = "Recipient entitlement month (1 through 12)." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_entitlement_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_entitlement_year.py new file mode 100644 index 00000000000..a4f82bc0ff1 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_entitlement_year.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_entitlement_year(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security survivors entitlement year" + documentation = "Recipient entitlement year; zero means unavailable." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_first_cola_month.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_first_cola_month.py new file mode 100644 index 00000000000..532c601f804 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_first_cola_month.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_first_cola_month(Variable): + value_type = int + default_value = 1 + entity = Person + definition_period = YEAR + label = "Social Security survivors first cola month" + documentation = "First exposed COLA payment month (1 through 12)." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_first_cola_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_first_cola_year.py new file mode 100644 index 00000000000..98d0ab59eb5 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_first_cola_year.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_first_cola_year(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security survivors first cola year" + documentation = "First exposed COLA payment year; zero requests the documented population proxy." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_joint_cola_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_joint_cola_factor.py new file mode 100644 index 00000000000..70056b3caa6 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_joint_cola_factor.py @@ -0,0 +1,17 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + reported_joint_cola_factor, +) + + +class social_security_survivors_joint_cola_factor(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Joint Social Security survivors COLA factor" + documentation = ( + "One payment-event numerator combines adjusted COLAs and withholding." + ) + + def formula(person, period, parameters): + return reported_joint_cola_factor(person, period, parameters, "survivors") diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_payable_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_payable_factor.py new file mode 100644 index 00000000000..0f8f0078d92 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_payable_factor.py @@ -0,0 +1,13 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_payable_factor(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Social Security survivors payable factor" + + def formula(person, period, parameters): + p = parameters(period).gov.contrib.ssa.benefit_levers.payable + factor = person("social_security_payable_factor", period) + return 1 if p.oasi_only and "survivors" == "disability" else factor diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_factor.py new file mode 100644 index 00000000000..8be462184ba --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_factor.py @@ -0,0 +1,14 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_pia_ratio_factor(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Social Security survivors PIA ratio factor" + documentation = ( + "Neutral extension point for a subsequent reported-benefit PIA reform." + ) + + def formula(person, period, parameters): + return 1 diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_reported.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_reported.py new file mode 100644 index 00000000000..e75781b950f --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_reported.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_reported(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Social Security survivors benefits, reported" + unit = USD + uprating = "calibration.gov.irs.soi.social_security" diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_source.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_source.py new file mode 100644 index 00000000000..40e8e669aa2 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_source.py @@ -0,0 +1,13 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, +) + + +class social_security_survivors_source(Variable): + value_type = Enum + possible_values = SocialSecurityBenefitSource + default_value = SocialSecurityBenefitSource.reported + entity = Person + definition_period = YEAR + label = "Social Security survivors benefit source" diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_eligibility_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_eligibility_year.py new file mode 100644 index 00000000000..889e023fefb --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_eligibility_year.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_worker_eligibility_year(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security survivors worker eligibility year" + documentation = "Explicit underlying-worker eligibility or computation year; zero means unavailable." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_id.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_id.py new file mode 100644 index 00000000000..50b868f4689 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_id.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_worker_id(Variable): + value_type = str + default_value = "" + entity = Person + definition_period = YEAR + label = "Social Security survivors worker id" + documentation = "Underlying-worker person identifier; requires an explicit computation year on that record." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_worker_computation_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_worker_computation_year.py new file mode 100644 index 00000000000..f5957bc0c97 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_worker_computation_year.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_worker_computation_year(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Supplied Social Security worker computation year" + documentation = "Authoritative underlying-record eligibility/computation year for linked COLA exposure, with zero indicating unavailable. Supply the legally applicable record basis; this input does not infer disability or deceased-worker computation history." From 2a81c22d12c681bf79cb32444d113e230ead7166 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 02:49:18 -0400 Subject: [PATCH 05/76] Isolate and extend Social Security compatibility coverage --- Makefile | 6 +++- .../test_social_security_current_law.py | 5 +++ .../test_social_security_inputs.py | 0 ...security_retirement_joint_cola_factor.yaml | 31 +++++++++++++++++++ 4 files changed, 41 insertions(+), 1 deletion(-) rename policyengine_us/tests/{core => microsimulation}/test_social_security_current_law.py (98%) rename policyengine_us/tests/{core => microsimulation}/test_social_security_inputs.py (100%) diff --git a/Makefile b/Makefile index a3dc2fcbe44..afcfb4cf12e 100644 --- a/Makefile +++ b/Makefile @@ -328,7 +328,11 @@ test-other-python-rest-remaining-heavy: test-other-python-rest-remaining: $(call rest_pytest,remaining,$(TESTS)/,$(REST_LISTED_TESTS) $(TESTS)/core $(TESTS)/policy) test-microsimulation: - pytest $(TESTS)/microsimulation --maxfail=0 + # Compatibility checks copy policy for differential/branch isolation. Keep + # their heaps separate on this existing runner (local differential peak 7 GB). + pytest $(TESTS)/microsimulation --maxfail=0 --ignore=$(TESTS)/microsimulation/test_social_security_inputs.py --ignore=$(TESTS)/microsimulation/test_social_security_current_law.py + pytest $(TESTS)/microsimulation/test_social_security_inputs.py --maxfail=0 + pytest $(TESTS)/microsimulation/test_social_security_current_law.py --maxfail=0 # Local convenience: run both halves back to back. test-other: test-other-python test-microsimulation test-policy-contrib-python: diff --git a/policyengine_us/tests/core/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py similarity index 98% rename from policyengine_us/tests/core/test_social_security_current_law.py rename to policyengine_us/tests/microsimulation/test_social_security_current_law.py index 45e562e3095..a3b78b0d97e 100644 --- a/policyengine_us/tests/core/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -339,6 +339,11 @@ def test_constructor_created_baseline_keeps_migrated_uprating( {"gov.contrib.ssa.benefit_levers.across_the_board.rate": {"2024": 0}} ) inputs = situation(years=(2024,)) + # Isolate this source migration from the existing wage/self-employment + # constructor moves, which Core's pre-wrapper baseline does not perform. + for person in inputs["people"].values(): + person.pop("employment_income") + person.pop("self_employment_income") original = reference_simulation( original_input_system, wrapper, situation=deepcopy(inputs) ) diff --git a/policyengine_us/tests/core/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py similarity index 100% rename from policyengine_us/tests/core/test_social_security_inputs.py rename to policyengine_us/tests/microsimulation/test_social_security_inputs.py diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_joint_cola_factor.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_joint_cola_factor.yaml index c6588bce46a..e8de4ba85d3 100644 --- a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_joint_cola_factor.yaml +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_joint_cola_factor.yaml @@ -116,3 +116,34 @@ social_security_retirement_reported: 12000 output: social_security_retirement_joint_cola_factor: 1.0014478764478765 + +- name: Adjacent eligibility-year sensitivity exposes a 63-year-old + period: 2027 + input: + age: 63 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + social_security_retirement_reported: 12000 + output: + social_security_retirement_joint_cola_factor: 1.0014478764478765 + +- name: A second configured CPI-E size uses the same formula + period: 2027 + input: + age: 70 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0030 + social_security_retirement_reported: 12000 + output: + social_security_retirement_joint_cola_factor: 1.0028957528957528 + +- name: An externally supplied alternative rate replaces the baseline numerator + period: 2027 + input: + age: 70 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.alternative_rate: 0.02 + social_security_retirement_first_cola_year: 2027 + social_security_retirement_reported: 12000 + output: + social_security_retirement_joint_cola_factor: 0.9845559845559845 From ee98c1d02fa11d459e991475d6ad5f41b27bc784 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 03:08:27 -0400 Subject: [PATCH 06/76] Preserve separate compatibility reports and reported-input uprating contracts --- Makefile | 12 ++++- .../code_health/test_rest_python_groups.py | 48 +++++++++++++++++++ .../parameters/test_uprating_extensions.py | 7 +-- .../ssa/test_trustees_core_thresholds.py | 26 ++++++++-- 4 files changed, 83 insertions(+), 10 deletions(-) diff --git a/Makefile b/Makefile index afcfb4cf12e..9a63a4a4b21 100644 --- a/Makefile +++ b/Makefile @@ -327,12 +327,20 @@ test-other-python-rest-remaining-heavy: $(call rest_pytest,remaining-heavy,$(REST_REMAINING_HEAVY_TESTS)) test-other-python-rest-remaining: $(call rest_pytest,remaining,$(TESTS)/,$(REST_LISTED_TESTS) $(TESTS)/core $(TESTS)/policy) +MICROSIMULATION_GROUPS := base social-security-inputs social-security-current-law test-microsimulation: + @failed=""; for group in $(MICROSIMULATION_GROUPS); do \ + $(MAKE) --no-print-directory test-microsimulation-$$group || failed="$$failed $$group"; \ + done; \ + if [ -n "$$failed" ]; then echo "Failed microsimulation groups:$$failed"; exit 1; fi +test-microsimulation-base: # Compatibility checks copy policy for differential/branch isolation. Keep # their heaps separate on this existing runner (local differential peak 7 GB). pytest $(TESTS)/microsimulation --maxfail=0 --ignore=$(TESTS)/microsimulation/test_social_security_inputs.py --ignore=$(TESTS)/microsimulation/test_social_security_current_law.py - pytest $(TESTS)/microsimulation/test_social_security_inputs.py --maxfail=0 - pytest $(TESTS)/microsimulation/test_social_security_current_law.py --maxfail=0 +test-microsimulation-social-security-inputs: + pytest $(TESTS)/microsimulation/test_social_security_inputs.py --maxfail=0 $(if $(findstring --junitxml,$(PYTEST_ADDOPTS)),--junitxml=simulation-social-security-inputs.xml) +test-microsimulation-social-security-current-law: + pytest $(TESTS)/microsimulation/test_social_security_current_law.py --maxfail=0 $(if $(findstring --junitxml,$(PYTEST_ADDOPTS)),--junitxml=simulation-social-security-current-law.xml) # Local convenience: run both halves back to back. test-other: test-other-python test-microsimulation test-policy-contrib-python: diff --git a/policyengine_us/tests/code_health/test_rest_python_groups.py b/policyengine_us/tests/code_health/test_rest_python_groups.py index 293b2a7f947..14a861176d9 100644 --- a/policyengine_us/tests/code_health/test_rest_python_groups.py +++ b/policyengine_us/tests/code_health/test_rest_python_groups.py @@ -577,6 +577,54 @@ def test_test_files_keep_distinct_import_names(reference): """ +@pytest.mark.parametrize("failed_index", [None, 0, 1, 2]) +def test_microsimulation_processes_keep_reports_and_run_after_failures( + tmp_path, failed_index +): + report_option = "--junitxml=simulation-microsimulation.xml" + commands = dry_run("test-microsimulation", f"PYTEST_ADDOPTS={report_option}") + assert len(commands) == 3 + reports = [] + groups = {} + for index, command in enumerate(commands): + args = [report_option, *command[1:]] + reports.append([a for a in args if a.startswith("--junitxml=")][-1]) + groups[str(index)] = collected_files( + [a for a in command[1:] if not a.startswith("--junitxml=")] + ) + reference = collected_files([str(TESTS / "microsimulation")]) + assert partition_problems(reference, groups) == [] + assert len(set(reports)) == 3 + assert all( + fnmatch(report.split("=", 1)[1], "simulation-*.xml") for report in reports + ) + + stub = tmp_path / "bin" / "pytest" + stub.parent.mkdir() + stub.write_text(STUB_PYTEST) + stub.chmod(0o755) + calls = tmp_path / "calls.txt" + outcomes = ["0"] * 3 + if failed_index is not None: + outcomes[failed_index] = "1" + env = {k: v for k, v in os.environ.items() if k not in PARENT_MAKE_VARIABLES} + env.update( + PATH=f"{stub.parent}{os.pathsep}{env['PATH']}", + PYTEST_ADDOPTS=report_option, + STUB_CALLS=str(calls), + STUB_OUTCOMES=" ".join(outcomes), + ) + result = subprocess.run( + ["make", "--no-print-directory", "test-microsimulation"], + cwd=REPO, + env=env, + capture_output=True, + text=True, + ) + assert calls.read_text().splitlines() == [" ".join(c[1:]) for c in commands] + assert (result.returncode != 0) == (failed_index is not None) + + @pytest.mark.parametrize("scenario", ["one-fails-one-killed", "all-fail", "all-pass"]) def test_every_group_runs_and_any_failure_fails_the_target(tmp_path, groups, scenario): outcomes = ["0"] * len(groups) diff --git a/policyengine_us/tests/policy/baseline/parameters/test_uprating_extensions.py b/policyengine_us/tests/policy/baseline/parameters/test_uprating_extensions.py index 6ff5f115dee..2acd8706939 100644 --- a/policyengine_us/tests/policy/baseline/parameters/test_uprating_extensions.py +++ b/policyengine_us/tests/policy/baseline/parameters/test_uprating_extensions.py @@ -130,9 +130,10 @@ def test_social_security_benefit_inputs_use_gross_benefit_uprater(): "social_security_survivors", "social_security_dependents", ]: - assert system.variables[variable_name].uprating == per_capita_path( - "calibration.gov.irs.soi.social_security" - ) + assert system.variables[ + f"{variable_name}_reported" + ].uprating == per_capita_path("calibration.gov.irs.soi.social_security") + assert system.variables[variable_name].uprating is None def test_cms_moop_per_capita_extends_to_2100(): diff --git a/policyengine_us/tests/policy/contrib/ssa/test_trustees_core_thresholds.py b/policyengine_us/tests/policy/contrib/ssa/test_trustees_core_thresholds.py index f9aa231fa5b..95a4bbe0322 100644 --- a/policyengine_us/tests/policy/contrib/ssa/test_trustees_core_thresholds.py +++ b/policyengine_us/tests/policy/contrib/ssa/test_trustees_core_thresholds.py @@ -86,25 +86,41 @@ def test_reform_only_changes_federal_irs_uprating_parameters(): assert changed_names -def test_social_security_components_use_aggregate_ss_uprater(): +def test_reported_social_security_components_use_aggregate_ss_uprater(): from policyengine_us.variables.gov.ssa.ss.social_security_dependents import ( social_security_dependents, ) + from policyengine_us.variables.gov.ssa.ss.social_security_dependents_reported import ( + social_security_dependents_reported, + ) from policyengine_us.variables.gov.ssa.ss.social_security_disability import ( social_security_disability, ) + from policyengine_us.variables.gov.ssa.ss.social_security_disability_reported import ( + social_security_disability_reported, + ) from policyengine_us.variables.gov.ssa.ss.social_security_retirement import ( social_security_retirement, ) + from policyengine_us.variables.gov.ssa.ss.social_security_retirement_reported import ( + social_security_retirement_reported, + ) from policyengine_us.variables.gov.ssa.ss.social_security_survivors import ( social_security_survivors, ) + from policyengine_us.variables.gov.ssa.ss.social_security_survivors_reported import ( + social_security_survivors_reported, + ) expected_uprater = "calibration.gov.irs.soi.social_security" - assert social_security_retirement.uprating == expected_uprater - assert social_security_disability.uprating == expected_uprater - assert social_security_dependents.uprating == expected_uprater - assert social_security_survivors.uprating == expected_uprater + for public, reported in ( + (social_security_retirement, social_security_retirement_reported), + (social_security_disability, social_security_disability_reported), + (social_security_dependents, social_security_dependents_reported), + (social_security_survivors, social_security_survivors_reported), + ): + assert reported.uprating == expected_uprater + assert getattr(public, "uprating", None) is None def test_reform_wage_indexes_core_tax_threshold_from_2035(): From 6101d97fdb6889f0d08eb92f93a3df89bef3c0df Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 03:25:43 -0400 Subject: [PATCH 07/76] Preserve overlapping Social Security sources and exact multi-year inputs --- docs/social-security-benefit-levers.md | 16 ++++ policyengine_us/system.py | 7 +- .../test_social_security_current_law.py | 34 ++++++++ .../test_social_security_inputs.py | 85 +++++++++++++++++++ .../tools/social_security_inputs.py | 40 +++++++-- .../tools/social_security_periods.py | 85 +++++++++++++++++++ .../ss/social_security_dependents_reported.py | 2 +- .../ss/social_security_retirement_reported.py | 2 +- .../ss/social_security_survivors_reported.py | 2 +- 9 files changed, 263 insertions(+), 10 deletions(-) create mode 100644 policyengine_us/tools/social_security_periods.py diff --git a/docs/social-security-benefit-levers.md b/docs/social-security-benefit-levers.md index a045fd076cf..5099f49934d 100644 --- a/docs/social-security-benefit-levers.md +++ b/docs/social-security-benefit-levers.md @@ -43,6 +43,22 @@ before final overrides. Core's legacy `extract_person` remains a one-period interface. Neither interface settles the separate API/application metadata and field-exposure contract. +Annual benefit inputs retain Core's period normalization. For example, supplying +`social_security_retirement: {"2024-07": 500}` through a PE-US constructor stores +500 under the rolling annual period `year:2024-07`, then migrates that exact +stored period to `_reported`. A July monthly output still follows Core's normal +calendar-year amount divided by twelve; it is not a new monthly entitlement +schedule. The complete source snapshot preserves the rolling annual key. + +PE-US wrappers also preserve an exact supplied multi-year annual total, such as +`year:2024:2`. They apply a benefit multiplier only when it is identical in every +calendar year covered by that total. Differing annual multipliers require +separately supplied annual source amounts, because the total provides no annual +allocation. Explicit public multi-year inputs remain final overrides, and +neutralization still takes precedence. Raw Core `SimulationBuilder` and YAML +should supply canonical annual `_reported` periods for lever calculations; the +additional exact multi-year reported-source adapter belongs to PE-US wrappers. + ## Parameters and composition All lever parameters live under `gov.contrib.ssa.benefit_levers`: diff --git a/policyengine_us/system.py b/policyengine_us/system.py index 25bf0c49b73..22c3d707221 100644 --- a/policyengine_us/system.py +++ b/policyengine_us/system.py @@ -37,6 +37,7 @@ from .tools.default_uprating import add_default_uprating from .tools.period_branch import delete_input_values from .tools.social_security_inputs import migrate_social_security_inputs +from .tools.social_security_periods import SocialSecurityPeriodMixin from .tools.section_911 import validate_section_911_batch_inputs from .tools.per_capita_uprating import ( add_per_capita_parameters_for_parameter_uprating, @@ -219,7 +220,7 @@ def _backfill_state_code_from_str(simulation): state_code_str.delete_arrays(known_period) -class Simulation(SPMSimulationMixin, CoreSimulation): +class Simulation(SocialSecurityPeriodMixin, SPMSimulationMixin, CoreSimulation): """ A simulation of the tax-benefit system for the United States, defined against the base simulation class in the -core package. @@ -422,7 +423,9 @@ def _is_hdfstore_format(file_path): return False -class Microsimulation(SPMSimulationMixin, CoreMicrosimulation): +class Microsimulation( + SocialSecurityPeriodMixin, SPMSimulationMixin, CoreMicrosimulation +): """ A microsimulation of the tax-benefit system for the United States, defined against the base microsimulation class in the -core package. diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index a3b78b0d97e..aa05ed93e99 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -354,3 +354,37 @@ def test_constructor_created_baseline_keeps_migrated_uprating( for public in SOCIAL_SECURITY_COMPONENTS: assert not current.baseline._get_set_input_periods(public) assert current.baseline._get_set_input_periods(f"{public}_reported") + + +@pytest.mark.parametrize("wrapper", [Simulation, HouseholdMicrosimulation]) +def test_monthly_input_normalization_and_exact_multi_year_reads_match_original( + original_input_system, wrapper +): + inputs = situation(include_benefits=False) + for person in inputs["people"].values(): + person["social_security_retirement"] = {2024: 12_000, "2024-07": 500} + person["social_security_disability"] = {"2024-07": 300} + person["social_security_survivors"] = {"year:2024:2": 2_000} + original = reference_simulation( + original_input_system, wrapper, situation=deepcopy(inputs) + ) + current = wrapper(situation=inputs, tax_benefit_system=system) + # Core's default FLOW input handler maps July to a rolling YEAR starting + # in July. It does not store an exact MONTH on these annual variables. + periods = current._get_set_input_periods("social_security_retirement_reported") + assert {str(period) for period in periods} == {"2024", "year:2024-07"} + assert all(period.unit == YEAR and period.size == 1 for period in periods) + requests = ("2024-07", "year:2024-07", 2024, "2024-06", "year:2024:2") + for public in SOCIAL_SECURITY_COMPONENTS[:2]: + assert_equal_outputs(current, original, requests[:-1], variables=(public,)) + assert_equal_outputs( + current, original, ("year:2024:2",), variables=("social_security_survivors",) + ) + reloaded = reload_source_inputs( + export_source_inputs(current), tax_benefit_system=system + ) + for public in SOCIAL_SECURITY_COMPONENTS[:2]: + assert_equal_outputs(reloaded, original, requests[:-1], variables=(public,)) + assert_equal_outputs( + reloaded, original, ("year:2024:2",), variables=("social_security_survivors",) + ) diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index ed0b5d03b2a..e39e8e2e3b2 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -235,6 +235,7 @@ def test_reported_twins_are_pure_and_original_dataset_names_keep_their_uprater() assert not original.is_input_variable() assert original.uprating is None assert reported.is_input_variable() + assert reported.label == original.label assert not reported.formulas and not reported.adds and not reported.subtracts assert reported.uprating.endswith("social_security_per_capita") assert ( @@ -303,3 +304,87 @@ def test_parameter_reform_registration_order_preserves_stable_public_formula(): first.tax_benefit_system.variables["social_security_retirement"].formulas == second.tax_benefit_system.variables["social_security_retirement"].formulas ) + + +@pytest.mark.parametrize("wrapper", WRAPPERS) +def test_exact_multi_year_sources_keep_keys_uniform_factors_and_overrides(wrapper): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + multi_year = "year:2024:2" + simulation = make_simulation(wrapper, {public: {multi_year: 24_000}}) + assert input_periods(simulation, public) == set() + assert input_periods(simulation, reported) == {multi_year} + assert amount(simulation, public, multi_year) == 24_000 + parent_keys = set(simulation._user_input_keys) + uniform = simulation.get_branch("uniform_multi_year") + for year in (2024, 2025): + uniform.set_input(FACTOR, year, [1.1]) + drop_inherited_values(uniform) + assert amount(uniform, public, multi_year) == pytest.approx(26_400) + assert simulation._user_input_keys == parent_keys + assert amount(simulation, public, multi_year) == 24_000 + reloaded = reload_source_inputs( + export_source_inputs(uniform), tax_benefit_system=system + ) + assert input_periods(reloaded, reported) == {multi_year} + assert amount(reloaded, public, multi_year) == pytest.approx(26_400) + uniform.set_input(public, multi_year, [0]) + clone = uniform.clone(clone_tax_benefit_system=False) + drop_inherited_values(clone) + assert amount(clone, public, multi_year) == 0 + override_reload = reload_source_inputs( + export_source_inputs(clone), tax_benefit_system=system + ) + assert amount(override_reload, public, multi_year) == 0 + assert amount(override_reload, reported, multi_year) == 24_000 + + +def test_exact_multi_year_both_names_and_neutralization_keep_core_precedence( + monkeypatch, +): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + multi_year = "year:2024:2" + simulation = make_simulation( + inputs={public: {multi_year: 123}, reported: {multi_year: 24_000}} + ) + simulation.set_input(FACTOR, 2024, [1.1]) + simulation.set_input(FACTOR, 2025, [0.8]) + assert amount(simulation, public, multi_year) == 123 + assert amount(simulation, reported, multi_year) == 24_000 + neutralized = make_simulation(inputs={reported: {multi_year: 24_000}}) + with monkeypatch.context() as context: + context.setattr(system.variables[public], "is_neutralized", True) + assert amount(neutralized, public, multi_year) == 0 + + +def test_exact_multi_year_source_rejects_differing_annual_factors_without_allocation(): + public = SOCIAL_SECURITY_COMPONENTS[0] + multi_year = "year:2024:2" + simulation = make_simulation(inputs={public: {multi_year: 24_000}}) + simulation.set_input(FACTOR, 2024, [1.1]) + simulation.set_input(FACTOR, 2025, [0.8]) + with pytest.raises(ValueError, match="separately supplied annual source amounts"): + simulation.calculate(public, multi_year) + assert input_periods(simulation, f"{public}_reported") == {multi_year} + + +def test_containing_period_migration_preserves_sources_and_inner_final_overrides(): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + simulation = make_simulation( + inputs={ + public: {2024: 123, "year:2024:10": 100_000, "year:2024:2": 24_000}, + reported: {2024: 12_000}, + } + ) + assert input_periods(simulation, public) == {"2024"} + assert input_periods(simulation, reported) == { + "2024", + "year:2024:2", + "year:2024:10", + } + assert amount(simulation, public, 2024) == 123 + assert amount(simulation, reported, 2024) == 12_000 + assert amount(simulation, public, "year:2024:2") == 24_000 + assert amount(simulation, reported, "year:2024:10") == 100_000 diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 59fca399937..b45cc864e98 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -33,16 +33,46 @@ def migrate_social_security_inputs(simulation, *, include_baseline=True): reported_name = f"{public_name}_reported" reported_periods = set(simulation._get_set_input_periods(reported_name)) holder = simulation.get_holder(public_name) + # Snapshot before deleting: deleting a multi-year period also deletes + # contained supplied periods, including a shorter legacy source whose + # lexical sort order happens to place it after the longer period. + sources = {} for period in simulation._get_set_input_periods(public_name): - if period in reported_periods: - continue - value = simulation._get_recorded_input_array(public_name, period) - if value is None: + if period not in reported_periods: + value = simulation._get_recorded_input_array(public_name, period) + if value is not None: + sources[period] = value.copy() + if not sources: + continue + retained_overrides = [] + for name, branch_name, period in tuple(simulation._user_input_keys): + if ( + name != public_name + or branch_name not in visible_branches + or period not in reported_periods + ): continue + for storage_name, storage in ( + ("memory", holder._memory_storage), + ("disk", holder._disk_storage), + ): + if holder._stores_user_input(period, branch_name, storage_name): + retained_overrides.append( + (branch_name, period, storage.get(period, branch_name).copy()) + ) + break + for period, value in sources.items(): simulation.set_input(reported_name, period, value.copy()) + for period in sources: for branch_name in visible_branches: delete_input_values(holder, period, branch_name) - migrated.add(public_name) + # Keep exact-period both-name overrides even inside a migrated longer + # source period. Restore each original branch slot through set_input, + # preserving its input classification and ordinary replay bookkeeping. + for branch_name, period, value in retained_overrides: + if not holder._stores_user_input(period, branch_name): + holder.set_input(period, value, branch_name) + migrated.add(public_name) # Core's legacy extraction reads this runtime list, which was assembled # before migration. Retain public names only where actual overrides remain. diff --git a/policyengine_us/tools/social_security_periods.py b/policyengine_us/tools/social_security_periods.py new file mode 100644 index 00000000000..a34838f6766 --- /dev/null +++ b/policyengine_us/tools/social_security_periods.py @@ -0,0 +1,85 @@ +"""Preserve exact multi-year benefit sources in PE-US wrapper simulations.""" + +import numpy as np +from policyengine_core.periods import YEAR +from policyengine_core.periods import period as to_period + +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, + SocialSecurityFund, +) +from policyengine_us.tools.social_security_inputs import SOCIAL_SECURITY_COMPONENTS + + +class SocialSecurityPeriodMixin: + def _calculate(self, variable_name, period=None): + if ( + variable_name not in SOCIAL_SECURITY_COMPONENTS + or period is None + or period.unit != YEAR + or period.size == 1 + ): + return super()._calculate(variable_name, period) + variable = self.tax_benefit_system.get_variable(variable_name) + # Core checks neutralization before exact public inputs and caches. + if variable.is_neutralized: + return super()._calculate(variable_name, period) + try: + abolished = self.tax_benefit_system.parameters(period).gov.abolitions[ + variable_name + ] + except Exception: + abolished = False + if abolished: + return super()._calculate(variable_name, period) + holder = self.get_holder(variable_name) + if holder.get_array(period, self.branch_name) is not None: + return super()._calculate(variable_name, period) + reported_name = f"{variable_name}_reported" + if period not in self._get_set_input_periods(reported_name): + return super()._calculate(variable_name, period) + + input_state = self._calculation_start() + amount = self.calculate(reported_name, period) + # Core accepts a stored multi-year YEAR input before checking period + # consistency. A uniform multiplier needs no annual allocation, while + # differing annual multipliers need separately supplied annual sources. + common_factor = None + for year_number in range(period.start.year, period.stop.year + 1): + year = to_period(year_number) + source = self.calculate(f"{variable_name}_source", year) + if np.any(source != SocialSecurityBenefitSource.reported): + raise NotImplementedError( + "Multi-year Social Security sources require the reported route; " + "calculated source modes are reserved" + ) + factor = np.ones_like(amount) + for name in ( + f"{variable_name}_joint_cola_factor", + f"{variable_name}_pia_ratio_factor", + "social_security_across_the_board_factor", + f"{variable_name}_payable_factor", + ): + factor = factor * self.calculate(name, year) + if ( + variable_name == "social_security_dependents" + and self.tax_benefit_system.parameters( + year + ).gov.contrib.ssa.benefit_levers.payable.oasi_only + and np.any(self.calculate("social_security_payable_factor", year) != 1) + ): + fund = self.calculate("social_security_dependents_fund", year) + if np.any((amount > 0) & (fund == SocialSecurityFund.unspecified)): + raise ValueError( + "OASI-only payable cuts require the dependent benefit fund" + ) + if common_factor is None: + common_factor = factor + elif not np.array_equal(common_factor, factor): + raise ValueError( + "Social Security levers with differing annual factors on a " + "multi-year reported input require separately supplied " + "annual source amounts" + ) + result = self._cast_formula_result(amount * common_factor, variable) + return self._cache_result(holder, result, period, input_state) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_reported.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_reported.py index a545607a6d6..bc04b45475b 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_reported.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_reported.py @@ -5,6 +5,6 @@ class social_security_dependents_reported(Variable): value_type = float entity = Person definition_period = YEAR - label = "Social Security dependents benefits, reported" + label = "Social Security dependents benefits" unit = USD uprating = "calibration.gov.irs.soi.social_security" diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_reported.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_reported.py index 5a89d1cbfc9..b8672e8bfc0 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_reported.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_reported.py @@ -5,6 +5,6 @@ class social_security_retirement_reported(Variable): value_type = float entity = Person definition_period = YEAR - label = "Social Security retirement benefits, reported" + label = "Social Security retirement benefits" unit = USD uprating = "calibration.gov.irs.soi.social_security" diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_reported.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_reported.py index e75781b950f..39fe64d8fd0 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_reported.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_reported.py @@ -5,6 +5,6 @@ class social_security_survivors_reported(Variable): value_type = float entity = Person definition_period = YEAR - label = "Social Security survivors benefits, reported" + label = "Social Security survivors benefits" unit = USD uprating = "calibration.gov.irs.soi.social_security" From b9c7191e51a11a6dfcc37f77c1ea2dbf0b0c16fb Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 03:36:21 -0400 Subject: [PATCH 08/76] Honor selected bases and zero rows for multi-year benefit sources --- docs/social-security-benefit-levers.md | 7 ++- .../test_social_security_inputs.py | 60 +++++++++++++++++++ .../tools/social_security_periods.py | 39 +++++++++++- 3 files changed, 101 insertions(+), 5 deletions(-) diff --git a/docs/social-security-benefit-levers.md b/docs/social-security-benefit-levers.md index 5099f49934d..95fc20d3e07 100644 --- a/docs/social-security-benefit-levers.md +++ b/docs/social-security-benefit-levers.md @@ -54,8 +54,11 @@ PE-US wrappers also preserve an exact supplied multi-year annual total, such as `year:2024:2`. They apply a benefit multiplier only when it is identical in every calendar year covered by that total. Differing annual multipliers require separately supplied annual source amounts, because the total provides no annual -allocation. Explicit public multi-year inputs remain final overrides, and -neutralization still takes precedence. Raw Core `SimulationBuilder` and YAML +allocation for nonzero rows. Zero rows remain zero. An exact supplied `_base` +amount takes precedence over `_reported`, and base neutralization or abolition +still applies; base overrides for only part of a multi-year source require +annual source amounts. Explicit public multi-year inputs remain final overrides, +and neutralization still takes precedence. Raw Core `SimulationBuilder` and YAML should supply canonical annual `_reported` periods for lever calculations; the additional exact multi-year reported-source adapter belongs to PE-US wrappers. diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index e39e8e2e3b2..2dd3d36c51d 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -388,3 +388,63 @@ def test_containing_period_migration_preserves_sources_and_inner_final_overrides assert amount(simulation, reported, 2024) == 12_000 assert amount(simulation, public, "year:2024:2") == 24_000 assert amount(simulation, reported, "year:2024:10") == 100_000 + + +def test_multi_year_selected_base_input_and_neutralization_take_precedence(monkeypatch): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported, base = f"{public}_reported", f"{public}_base" + multi_year = "year:2024:2" + simulation = make_simulation( + inputs={reported: {multi_year: 24_000}, base: {multi_year: 6_000}} + ) + for year in (2024, 2025): + simulation.set_input(FACTOR, year, [1.1]) + assert amount(simulation, public, multi_year) == pytest.approx(6_600) + reloaded = reload_source_inputs( + export_source_inputs(simulation), tax_benefit_system=system + ) + assert amount(reloaded, public, multi_year) == pytest.approx(6_600) + + base_only = make_simulation(inputs={base: {multi_year: 6_000}}) + assert amount(base_only, public, multi_year) == 6_000 + neutralized = make_simulation( + inputs={reported: {multi_year: 24_000}, base: {multi_year: 6_000}} + ) + neutralized.set_input(FACTOR, 2024, [1.1]) + neutralized.set_input(FACTOR, 2025, [0.8]) + with monkeypatch.context() as context: + context.setattr(system.variables[base], "is_neutralized", True) + assert amount(neutralized, public, multi_year) == 0 + + +def test_multi_year_reported_total_rejects_partial_base_overrides(): + public = SOCIAL_SECURITY_COMPONENTS[0] + simulation = make_simulation( + inputs={ + f"{public}_reported": {"year:2024:2": 24_000}, + f"{public}_base": {2024: 6_000}, + } + ) + with pytest.raises(ValueError, match="base overrides for other periods"): + simulation.calculate(public, "year:2024:2") + + +def test_multi_year_zero_rows_do_not_require_annual_allocation(): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + multi_year = "year:2024:2" + inputs = { + "people": { + "positive": {reported: {multi_year: 24_000}}, + "zero": {reported: {multi_year: 0}}, + }, + "households": {"household": {"members": ["positive", "zero"]}}, + } + simulation = Simulation(situation=inputs, tax_benefit_system=system) + simulation.set_input(FACTOR, 2024, [1.1, 1.1]) + simulation.set_input(FACTOR, 2025, [1.1, 0.8]) + np.testing.assert_allclose(simulation.calculate(public, multi_year), [26_400, 0]) + all_zero = simulation.get_branch("all_zero_multi_year") + all_zero.set_input(reported, multi_year, [0, 0]) + drop_inherited_values(all_zero) + np.testing.assert_array_equal(all_zero.calculate(public, multi_year), [0, 0]) diff --git a/policyengine_us/tools/social_security_periods.py b/policyengine_us/tools/social_security_periods.py index a34838f6766..fbfa06b458e 100644 --- a/policyengine_us/tools/social_security_periods.py +++ b/policyengine_us/tools/social_security_periods.py @@ -36,11 +36,36 @@ def _calculate(self, variable_name, period=None): if holder.get_array(period, self.branch_name) is not None: return super()._calculate(variable_name, period) reported_name = f"{variable_name}_reported" - if period not in self._get_set_input_periods(reported_name): + base_name = f"{variable_name}_base" + base_periods = self._get_set_input_periods(base_name) + has_exact_base = period in base_periods + if not has_exact_base and period not in self._get_set_input_periods( + reported_name + ): return super()._calculate(variable_name, period) input_state = self._calculation_start() - amount = self.calculate(reported_name, period) + base_variable = self.tax_benefit_system.get_variable(base_name) + try: + base_abolished = self.tax_benefit_system.parameters(period).gov.abolitions[ + base_name + ] + except Exception: + base_abolished = False + if has_exact_base or base_variable.is_neutralized or base_abolished: + # Let Core preserve exact base-input and neutralization precedence. + amount = self.calculate(base_name, period) + else: + if any( + known.start <= period.stop and known.stop >= period.start + for known in base_periods + ): + raise ValueError( + "A multi-year reported input with base overrides for other " + "periods requires separately supplied annual source amounts" + ) + amount = self.calculate(reported_name, period) + nonzero = amount != 0 # Core accepts a stored multi-year YEAR input before checking period # consistency. A uniform multiplier needs no annual allocation, while # differing annual multipliers need separately supplied annual sources. @@ -61,6 +86,14 @@ def _calculate(self, variable_name, period=None): f"{variable_name}_payable_factor", ): factor = factor * self.calculate(name, year) + try: + annual_base_abolished = self.tax_benefit_system.parameters( + year + ).gov.abolitions[base_name] + except Exception: + annual_base_abolished = False + if annual_base_abolished: + factor = np.zeros_like(factor) if ( variable_name == "social_security_dependents" and self.tax_benefit_system.parameters( @@ -75,7 +108,7 @@ def _calculate(self, variable_name, period=None): ) if common_factor is None: common_factor = factor - elif not np.array_equal(common_factor, factor): + elif not np.array_equal(common_factor[nonzero], factor[nonzero]): raise ValueError( "Social Security levers with differing annual factors on a " "multi-year reported input require separately supplied " From 89edc6c260e828dc9aa1511650e92724a89f428a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 03:37:20 -0400 Subject: [PATCH 09/76] Share identical COLA policy fixtures across exposure cases --- ...security_retirement_joint_cola_factor.yaml | 116 ++++++++---------- 1 file changed, 54 insertions(+), 62 deletions(-) diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_joint_cola_factor.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_joint_cola_factor.yaml index e8de4ba85d3..d5e07de6e72 100644 --- a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_joint_cola_factor.yaml +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_joint_cola_factor.yaml @@ -1,26 +1,62 @@ -- name: CPI-E precision survives the statutory baseline rounding +- name: A 2027 CPI-E change preserves precision and retirement exposure boundaries period: 2027 - absolute_error_margin: 0.01 + absolute_error_margin: 0.001 input: - age: 70 gov.contrib.ssa.benefit_levers.cola.start_year: 2027 gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 - social_security_retirement_reported: 12000 + people: + precision_70: + # CPI-E precision survives the statutory baseline rounding. + age: 70 + social_security_retirement_reported: 12000 + newly_eligible_62: + # Newly eligible retirement has no same-year COLA exposure. + age: 62 + social_security_retirement_reported: 12000 + late_claim_67: + # Retirement COLAs accrue before a later claiming age. + age: 67 + ss_claiming_age: 67 + social_security_retirement_reported: 12000 + supplied_first_cola_75: + # Supplied first-COLA date overrides an older-age proxy. + age: 75 + social_security_retirement_first_cola_year: 2028 + social_security_retirement_reported: 12000 + late_entitlement_70: + # A retirement claim date does not erase eligibility COLAs. + age: 70 + social_security_retirement_entitlement_year: 2027 + social_security_retirement_reported: 12000 + adjacent_eligibility_63: + # Adjacent eligibility-year sensitivity exposes a 63-year-old. + age: 63 + social_security_retirement_reported: 12000 + reported_zero_70: + # A reported zero remains zero despite positive COLA exposure. + age: 70 + social_security_retirement_reported: 0 output: - # January 2027 baseline is 3.6%; 12000 * 1.0375 / 1.036. - social_security_retirement_joint_cola_factor: 1.0014478764478765 - social_security_retirement: 12017.374517374517 - -- name: Newly eligible retirement has no same-year COLA exposure - period: 2027 - input: - age: 62 - gov.contrib.ssa.benefit_levers.cola.start_year: 2027 - gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 - social_security_retirement_reported: 12000 - output: - social_security_retirement_joint_cola_factor: 1 - social_security_retirement: 12000 + people: + precision_70: + # January 2027 baseline is 3.6%; 12000 * 1.0375 / 1.036. + social_security_retirement_joint_cola_factor: 1.0014478764478765 + social_security_retirement: 12017.374517374517 + newly_eligible_62: + social_security_retirement_joint_cola_factor: 1 + social_security_retirement: 12000 + late_claim_67: + social_security_retirement_joint_cola_factor: 1.0014478764478765 + supplied_first_cola_75: + social_security_retirement_joint_cola_factor: 1 + social_security_retirement: 12000 + late_entitlement_70: + social_security_retirement_joint_cola_factor: 1.0014478764478765 + adjacent_eligibility_63: + social_security_retirement_joint_cola_factor: 1.0014478764478765 + reported_zero_70: + social_security_retirement_joint_cola_factor: 1.0014478764478765 + social_security_retirement: 0 - name: Static age cross-sections infer eligibility separately in each year period: 2035 @@ -35,29 +71,6 @@ social_security_retirement_joint_cola_factor: 1 social_security_retirement: 12000 -- name: Retirement COLAs accrue before a later claiming age - period: 2027 - input: - age: 67 - ss_claiming_age: 67 - gov.contrib.ssa.benefit_levers.cola.start_year: 2027 - gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 - social_security_retirement_reported: 12000 - output: - social_security_retirement_joint_cola_factor: 1.0014478764478765 - -- name: Supplied first-COLA date overrides an older-age proxy - period: 2027 - input: - age: 75 - social_security_retirement_first_cola_year: 2028 - gov.contrib.ssa.benefit_levers.cola.start_year: 2027 - gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 - social_security_retirement_reported: 12000 - output: - social_security_retirement_joint_cola_factor: 1 - social_security_retirement: 12000 - - name: Withheld COLA uses one numerator even when CPI-E is active period: 2027 absolute_error_margin: 0.01 @@ -106,27 +119,6 @@ # Statutory baseline 339.431 / 332.241 - 1 rounds to 2.2%. social_security_retirement_joint_cola_factor: 1.0014677103718199 -- name: A retirement claim date does not erase eligibility COLAs - period: 2027 - input: - age: 70 - social_security_retirement_entitlement_year: 2027 - gov.contrib.ssa.benefit_levers.cola.start_year: 2027 - gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 - social_security_retirement_reported: 12000 - output: - social_security_retirement_joint_cola_factor: 1.0014478764478765 - -- name: Adjacent eligibility-year sensitivity exposes a 63-year-old - period: 2027 - input: - age: 63 - gov.contrib.ssa.benefit_levers.cola.start_year: 2027 - gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 - social_security_retirement_reported: 12000 - output: - social_security_retirement_joint_cola_factor: 1.0014478764478765 - - name: A second configured CPI-E size uses the same formula period: 2027 input: From 8401cd89c6abd68ae4b39c3508171159c05de26f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 03:51:23 -0400 Subject: [PATCH 10/76] Treat Social Security benefit multipliers as dimensionless levels --- .../test_social_security_inputs.py | 24 +++++++++++++++++++ ...social_security_across_the_board_factor.py | 1 + ...l_security_dependents_joint_cola_factor.py | 1 + ...cial_security_dependents_payable_factor.py | 1 + ...al_security_dependents_pia_ratio_factor.py | 1 + ...l_security_disability_joint_cola_factor.py | 1 + ...cial_security_disability_payable_factor.py | 1 + ...al_security_disability_pia_ratio_factor.py | 1 + .../ssa/ss/social_security_payable_factor.py | 1 + ...l_security_retirement_joint_cola_factor.py | 1 + ...cial_security_retirement_payable_factor.py | 1 + ...al_security_retirement_pia_ratio_factor.py | 1 + ...al_security_survivors_joint_cola_factor.py | 1 + ...ocial_security_survivors_payable_factor.py | 1 + ...ial_security_survivors_pia_ratio_factor.py | 1 + 15 files changed, 38 insertions(+) diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index 2dd3d36c51d..a45d5e07cfd 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -448,3 +448,27 @@ def test_multi_year_zero_rows_do_not_require_annual_allocation(): all_zero.set_input(reported, multi_year, [0, 0]) drop_inherited_values(all_zero) np.testing.assert_array_equal(all_zero.calculate(public, multi_year), [0, 0]) + + +def test_dimensionless_benefit_factors_keep_annual_levels_on_monthly_requests(): + factor_names = ( + "social_security_across_the_board_factor", + "social_security_payable_factor", + *( + f"{component}_{suffix}" + for component in SOCIAL_SECURITY_COMPONENTS + for suffix in ("joint_cola_factor", "pia_ratio_factor", "payable_factor") + ), + ) + simulation = make_simulation() + for name in factor_names: + assert system.variables[name].unit == "/1" + assert amount(simulation, name, "2024-07") == 1 + assert amount(simulation, name, 2024) == 1 + supplied = simulation.get_branch("supplied_dimensionless_factors") + for name in factor_names: + supplied.set_input(name, 2024, [1.2]) + drop_inherited_values(supplied) + for name in factor_names: + assert amount(supplied, name, "2024-07") == pytest.approx(1.2) + assert amount(supplied, name, 2024) == pytest.approx(1.2) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_across_the_board_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_across_the_board_factor.py index d5c65cc9b14..ea4950d9cc7 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_across_the_board_factor.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_across_the_board_factor.py @@ -5,6 +5,7 @@ class social_security_across_the_board_factor(Variable): value_type = float entity = Person definition_period = YEAR + unit = "/1" label = "Social Security across-the-board factor" def formula(person, period, parameters): diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_joint_cola_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_joint_cola_factor.py index 64539ed253b..e12dde9a5fa 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_joint_cola_factor.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_joint_cola_factor.py @@ -8,6 +8,7 @@ class social_security_dependents_joint_cola_factor(Variable): value_type = float entity = Person definition_period = YEAR + unit = "/1" label = "Joint Social Security dependents COLA factor" documentation = ( "One payment-event numerator combines adjusted COLAs and withholding." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_payable_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_payable_factor.py index b1beac8c468..03739259995 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_payable_factor.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_payable_factor.py @@ -6,6 +6,7 @@ class social_security_dependents_payable_factor(Variable): value_type = float entity = Person definition_period = YEAR + unit = "/1" label = "Social Security dependents payable factor" def formula(person, period, parameters): diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_factor.py index 99543e62b56..034b50023c9 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_factor.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_factor.py @@ -5,6 +5,7 @@ class social_security_dependents_pia_ratio_factor(Variable): value_type = float entity = Person definition_period = YEAR + unit = "/1" label = "Social Security dependents PIA ratio factor" documentation = ( "Neutral extension point for a subsequent reported-benefit PIA reform." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_joint_cola_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_joint_cola_factor.py index e2a73b44adf..d088d10392b 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_disability_joint_cola_factor.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_joint_cola_factor.py @@ -8,6 +8,7 @@ class social_security_disability_joint_cola_factor(Variable): value_type = float entity = Person definition_period = YEAR + unit = "/1" label = "Joint Social Security disability COLA factor" documentation = ( "One payment-event numerator combines adjusted COLAs and withholding." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_payable_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_payable_factor.py index efbfaafb98e..947c742323c 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_disability_payable_factor.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_payable_factor.py @@ -5,6 +5,7 @@ class social_security_disability_payable_factor(Variable): value_type = float entity = Person definition_period = YEAR + unit = "/1" label = "Social Security disability payable factor" def formula(person, period, parameters): diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_factor.py index 3d641e7cf60..e807b247634 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_factor.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_factor.py @@ -5,6 +5,7 @@ class social_security_disability_pia_ratio_factor(Variable): value_type = float entity = Person definition_period = YEAR + unit = "/1" label = "Social Security disability PIA ratio factor" documentation = ( "Neutral extension point for a subsequent reported-benefit PIA reform." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_payable_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_payable_factor.py index 9f653f8f065..c9834a8d7ee 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_payable_factor.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_payable_factor.py @@ -5,6 +5,7 @@ class social_security_payable_factor(Variable): value_type = float entity = Person definition_period = YEAR + unit = "/1" label = "Social Security supplied payable-share factor" documentation = "Externally supplied annual share, including a supplier-blended depletion year; no depletion projection is inferred." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_joint_cola_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_joint_cola_factor.py index 55214555c25..668e6957e8e 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_joint_cola_factor.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_joint_cola_factor.py @@ -8,6 +8,7 @@ class social_security_retirement_joint_cola_factor(Variable): value_type = float entity = Person definition_period = YEAR + unit = "/1" label = "Joint Social Security retirement COLA factor" documentation = ( "One payment-event numerator combines adjusted COLAs and withholding." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_payable_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_payable_factor.py index f88697a58ba..43f8dcbce80 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_payable_factor.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_payable_factor.py @@ -5,6 +5,7 @@ class social_security_retirement_payable_factor(Variable): value_type = float entity = Person definition_period = YEAR + unit = "/1" label = "Social Security retirement payable factor" def formula(person, period, parameters): diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_factor.py index fd3609292ac..50375513765 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_factor.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_factor.py @@ -5,6 +5,7 @@ class social_security_retirement_pia_ratio_factor(Variable): value_type = float entity = Person definition_period = YEAR + unit = "/1" label = "Social Security retirement PIA ratio factor" documentation = ( "Neutral extension point for a subsequent reported-benefit PIA reform." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_joint_cola_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_joint_cola_factor.py index 70056b3caa6..b103db77f4c 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_joint_cola_factor.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_joint_cola_factor.py @@ -8,6 +8,7 @@ class social_security_survivors_joint_cola_factor(Variable): value_type = float entity = Person definition_period = YEAR + unit = "/1" label = "Joint Social Security survivors COLA factor" documentation = ( "One payment-event numerator combines adjusted COLAs and withholding." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_payable_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_payable_factor.py index 0f8f0078d92..44ac4c71a4a 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_payable_factor.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_payable_factor.py @@ -5,6 +5,7 @@ class social_security_survivors_payable_factor(Variable): value_type = float entity = Person definition_period = YEAR + unit = "/1" label = "Social Security survivors payable factor" def formula(person, period, parameters): diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_factor.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_factor.py index 8be462184ba..40be5f4d21d 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_factor.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_factor.py @@ -5,6 +5,7 @@ class social_security_survivors_pia_ratio_factor(Variable): value_type = float entity = Person definition_period = YEAR + unit = "/1" label = "Social Security survivors PIA ratio factor" documentation = ( "Neutral extension point for a subsequent reported-benefit PIA reform." From fa7ddfad8fcdc50b289a8bc786cef2fa49e620ed Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 04:08:38 -0400 Subject: [PATCH 11/76] Run Social Security lever YAML serially on the existing CI runner --- Makefile | 8 +++-- .../code_health/test_make_test_matches_ci.py | 30 +++++++++++++++++++ 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/Makefile b/Makefile index 9a63a4a4b21..e25f7d79e56 100644 --- a/Makefile +++ b/Makefile @@ -166,10 +166,12 @@ test-yaml-structural-heavy-shard-4: test-yaml-structural-other: # Per-subdir so every remaining contrib folder runs in its own subprocess, # instead of stacking ~20 light files into one ~13-min catch-all batch that - # risked the 30-min per-batch timeout on slow runners. ssa is excluded (it - # has no YAML tests — only a pytest .py run elsewhere). --workers 1 so a - # single ~1-5 GB peak is resident at a time, leaving headroom for growth. + # risked the 30-min per-batch timeout on slow runners. SSA runs separately + # per-file below because its reform combinations retain large model copies. $(BATCH) $(TESTS)/policy/contrib --exclude states,ctc,ubi_center,federal,harris,treasury,crfb,congress,refundable_credit_conversion,ssa --mode per-subdir --workers 1 + # SSA's COLA file measured about 11 GB locally. Override CI's worker setting + # as well as the CLI default so each file exits before the next one starts. + POLICYENGINE_BATCH_WORKERS=1 $(BATCH) $(TESTS)/policy/contrib/ssa --mode per-file --workers 1 test-yaml-structural-other-shard-2a: # ctc is microsim-heavy: per-file isolation frees each ~5 GB peak between # files, and --workers 1 keeps only one peak resident at a time. diff --git a/policyengine_us/tests/code_health/test_make_test_matches_ci.py b/policyengine_us/tests/code_health/test_make_test_matches_ci.py index 532945e5cf9..c06788638a5 100644 --- a/policyengine_us/tests/code_health/test_make_test_matches_ci.py +++ b/policyengine_us/tests/code_health/test_make_test_matches_ci.py @@ -8,6 +8,7 @@ """ import re +import shlex import shutil import subprocess from functools import lru_cache @@ -103,3 +104,32 @@ def test_make_test_never_runs_a_whole_tree_in_one_process(): "run YAML suites through policyengine_us/tests/test_batched.py, " f"not one policyengine-core process: {single_process}" ) + + +def test_contrib_ssa_yaml_has_a_dedicated_single_file_serial_route(): + """Local/CI equality alone missed SSA when both routes excluded its YAML. + + Its large reform combinations need separate files and an effective + one-worker limit even when CI exports POLICYENGINE_BATCH_WORKERS=2. + """ + root = "policyengine_us/tests/policy/contrib" + ssa = f"{root}/ssa" + assert any((REPO_ROOT / ssa).rglob("*.yaml")), "SSA YAML coverage is absent" + commands = make_commands("test-yaml-structural-other") + batcher = "policyengine_us/tests/test_batched.py" + batch_commands = [ + shlex.split(command) for command in commands if batcher in command + ] + dedicated = [arguments for arguments in batch_commands if ssa in arguments] + assert len(dedicated) == 1, "SSA must have one dedicated YAML invocation" + (arguments,) = dedicated + assert arguments[arguments.index("--mode") + 1] == "per-file" + assert arguments[arguments.index("--workers") + 1] == "1" + assert "POLICYENGINE_BATCH_WORKERS=1" in arguments + + generic = [arguments for arguments in batch_commands if root in arguments] + assert len(generic) == 1 + (arguments,) = generic + excluded = arguments[arguments.index("--exclude") + 1].split(",") + assert "ssa" in excluded, "SSA YAML must not also enter a shared subdir batch" + assert commands <= make_commands("test"), "The full suite must retain SSA's route" From 561a66494b31584b80c7283902e1bc0a26b822df Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 05:06:55 -0400 Subject: [PATCH 12/76] Clarify the Core public retirement override fixture --- .../gov/ssa/social_security/social_security_retirement.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/policyengine_us/tests/policy/baseline/gov/ssa/social_security/social_security_retirement.yaml b/policyengine_us/tests/policy/baseline/gov/ssa/social_security/social_security_retirement.yaml index 8c3799ea037..f551b4ca2ea 100644 --- a/policyengine_us/tests/policy/baseline/gov/ssa/social_security/social_security_retirement.yaml +++ b/policyengine_us/tests/policy/baseline/gov/ssa/social_security/social_security_retirement.yaml @@ -1,4 +1,4 @@ -- name: Social Security retirement benefits are an input variable +- name: Core public retirement inputs remain final overrides period: 2024 input: social_security_retirement: 18_000 From fbca48fcd0dfc0a9fbb08f606bc5703eac6ba266 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 05:16:09 -0400 Subject: [PATCH 13/76] Verify original Social Security dataset columns extend once --- .../test_social_security_inputs.py | 63 ++++++++++++++++++- 1 file changed, 61 insertions(+), 2 deletions(-) diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index a45d5e07cfd..b3930086e49 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -2,7 +2,7 @@ These cases need Python because YAML bypasses PE-US constructors and cannot exercise cloning, stored-input bookkeeping or source export. They reuse the -read-only country model and use one person, without loading population data. +read-only country model and tiny populations, without loading external data. """ import json @@ -13,7 +13,11 @@ from policyengine_core.simulations import SimulationBuilder from policyengine_us import Microsimulation, Simulation -from policyengine_us.data.economic_assumptions import MICRODATA_UPRATING_OVERRIDES +from policyengine_us.data.dataset_schema import USSingleYearDataset +from policyengine_us.data.economic_assumptions import ( + MICRODATA_UPRATING_OVERRIDES, + extend_single_year_dataset, +) from policyengine_us.system import system from policyengine_us.tools.period_branch import drop_inherited_values from policyengine_us.tools.social_security_inputs import ( @@ -472,3 +476,58 @@ def test_dimensionless_benefit_factors_keep_annual_levels_on_monthly_requests(): for name in factor_names: assert amount(supplied, name, "2024-07") == pytest.approx(1.2) assert amount(supplied, name, 2024) == pytest.approx(1.2) + + +def test_original_component_dataset_columns_extend_once_by_per_capita_growth(): + base_amounts = dict(zip(SOCIAL_SECURITY_COMPONENTS, (24_000, 6_000, 3_000, 1_200))) + group_entities = ("household", "tax_unit", "spm_unit", "family", "marital_unit") + single = USSingleYearDataset( + person=pd.DataFrame( + { + "person_id": [1], + "age": [70], + **{f"person_{entity}_id": [1] for entity in group_entities}, + **{name: [value] for name, value in base_amounts.items()}, + } + ), + household=pd.DataFrame({"household_id": [1], "household_weight": [1.0]}), + tax_unit=pd.DataFrame({"tax_unit_id": [1]}), + spm_unit=pd.DataFrame({"spm_unit_id": [1]}), + family=pd.DataFrame({"family_id": [1]}), + marital_unit=pd.DataFrame({"marital_unit_id": [1]}), + time_period=2024, + ) + # Independent expectation: national-total growth divided by population + # growth exactly once, without using the dataset uprater's resolver. + base_parameters = system.parameters("2024-01-01") + projected_parameters = system.parameters("2025-01-01") + national_growth = ( + projected_parameters.calibration.gov.irs.soi.social_security + / base_parameters.calibration.gov.irs.soi.social_security + ) + population_growth = ( + projected_parameters.calibration.gov.census.populations.total + / base_parameters.calibration.gov.census.populations.total + ) + per_person_growth = national_growth / population_growth + extended = extend_single_year_dataset(single, end_year=2025, system=system) + for public, base in base_amounts.items(): + assert single.person[public].iloc[0] == base + assert extended[2024].person[public].iloc[0] == base + assert extended[2025].person[public].iloc[0] == pytest.approx( + base * per_person_growth, rel=1e-12 + ) + + simulation = Microsimulation(dataset=extended, tax_benefit_system=system) + for public, base in base_amounts.items(): + reported = f"{public}_reported" + assert input_periods(simulation, public) == set() + assert input_periods(simulation, reported) == {"2024", "2025"} + for year, expected in ((2024, base), (2025, base * per_person_growth)): + assert amount(simulation, public, f"{year}-07") == pytest.approx( + expected / 12, rel=1e-6 + ) + assert amount(simulation, public, year) == pytest.approx(expected, rel=1e-6) + assert amount(simulation, reported, year) == pytest.approx( + expected, rel=1e-6 + ) From 460f573eba424d56d96f5845f64183425b368e08 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 05:19:58 -0400 Subject: [PATCH 14/76] Preserve benefit metadata and explicit multi-year test weights --- .../microsimulation/test_social_security_current_law.py | 4 ++++ .../tests/microsimulation/test_social_security_inputs.py | 6 ++++++ .../variables/gov/ssa/ss/social_security_disability.py | 2 +- 3 files changed, 11 insertions(+), 1 deletion(-) diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index aa05ed93e99..c015b5e596e 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -365,6 +365,10 @@ def test_monthly_input_normalization_and_exact_multi_year_reads_match_original( person["social_security_retirement"] = {2024: 12_000, "2024-07": 500} person["social_security_disability"] = {"2024-07": 300} person["social_security_survivors"] = {"year:2024:2": 2_000} + for household in inputs["households"].values(): + # Main's weighted API also needs an exact weight for this multi-year + # request; source migration does not change Core's weight semantics. + household["household_weight"] = {"year:2024:2": 1} original = reference_simulation( original_input_system, wrapper, situation=deepcopy(inputs) ) diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index b3930086e49..786b14f36bd 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -316,6 +316,9 @@ def test_exact_multi_year_sources_keep_keys_uniform_factors_and_overrides(wrappe reported = f"{public}_reported" multi_year = "year:2024:2" simulation = make_simulation(wrapper, {public: {multi_year: 24_000}}) + # Core's weighted Microsimulation API needs a weight for this exact + # noncanonical request, just as it does for an exact public input on main. + simulation.set_input("household_weight", multi_year, [1]) assert input_periods(simulation, public) == set() assert input_periods(simulation, reported) == {multi_year} assert amount(simulation, public, multi_year) == 24_000 @@ -331,16 +334,19 @@ def test_exact_multi_year_sources_keep_keys_uniform_factors_and_overrides(wrappe export_source_inputs(uniform), tax_benefit_system=system ) assert input_periods(reloaded, reported) == {multi_year} + assert input_periods(reloaded, "household_weight") == {multi_year} assert amount(reloaded, public, multi_year) == pytest.approx(26_400) uniform.set_input(public, multi_year, [0]) clone = uniform.clone(clone_tax_benefit_system=False) drop_inherited_values(clone) + assert input_periods(clone, "household_weight") == {multi_year} assert amount(clone, public, multi_year) == 0 override_reload = reload_source_inputs( export_source_inputs(clone), tax_benefit_system=system ) assert amount(override_reload, public, multi_year) == 0 assert amount(override_reload, reported, multi_year) == 24_000 + assert input_periods(override_reload, "household_weight") == {multi_year} def test_exact_multi_year_both_names_and_neutralization_keep_core_precedence( diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability.py index 540f39139a5..b9987fb0bf1 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_disability.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability.py @@ -8,7 +8,7 @@ class social_security_disability(Variable): value_type = float entity = Person definition_period = YEAR - label = "Social Security disability benefits" + label = "Social Security disability benefits (SSDI)" unit = USD documentation = "Selected source with applicable benefit levers; an explicit public input is a final override." From 4789c5408b13cb1b8e32f60f1c7d165718d98aa6 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 05:23:54 -0400 Subject: [PATCH 15/76] Check declared Social Security uprating metadata --- .../tests/policy/contrib/ssa/test_trustees_core_thresholds.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/policyengine_us/tests/policy/contrib/ssa/test_trustees_core_thresholds.py b/policyengine_us/tests/policy/contrib/ssa/test_trustees_core_thresholds.py index 95a4bbe0322..dd0f0d5942a 100644 --- a/policyengine_us/tests/policy/contrib/ssa/test_trustees_core_thresholds.py +++ b/policyengine_us/tests/policy/contrib/ssa/test_trustees_core_thresholds.py @@ -120,7 +120,7 @@ def test_reported_social_security_components_use_aggregate_ss_uprater(): (social_security_survivors, social_security_survivors_reported), ): assert reported.uprating == expected_uprater - assert getattr(public, "uprating", None) is None + assert "uprating" not in public.__dict__ def test_reform_wage_indexes_core_tax_threshold_from_2035(): From 99d738450c26739c42977e8f9f3b6d376c3b76a5 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 05:28:26 -0400 Subject: [PATCH 16/76] Use raw calculation arrays inside multi-year benefit factors --- docs/social-security-benefit-levers.md | 2 ++ .../test_social_security_inputs.py | 29 +++++++++++++++++++ .../tools/social_security_periods.py | 20 +++++++++---- 3 files changed, 45 insertions(+), 6 deletions(-) diff --git a/docs/social-security-benefit-levers.md b/docs/social-security-benefit-levers.md index 95fc20d3e07..9a7e974e9d6 100644 --- a/docs/social-security-benefit-levers.md +++ b/docs/social-security-benefit-levers.md @@ -61,6 +61,8 @@ annual source amounts. Explicit public multi-year inputs remain final overrides, and neutralization still takes precedence. Raw Core `SimulationBuilder` and YAML should supply canonical annual `_reported` periods for lever calculations; the additional exact multi-year reported-source adapter belongs to PE-US wrappers. +A weighted `Microsimulation` request also needs an explicit weight for that exact +multi-year period, matching Core. ## Parameters and composition diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index 786b14f36bd..4059a7c773c 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -537,3 +537,32 @@ def test_original_component_dataset_columns_extend_once_by_per_capita_growth(): assert amount(simulation, reported, year) == pytest.approx( expected, rel=1e-6 ) + + +def test_multi_year_microsimulation_uses_raw_source_and_fund_enums(): + from policyengine_core.reforms import Reform + + public = "social_security_dependents" + multi_year = "year:2024:2" + reform = Reform.from_dict( + { + "gov.contrib.ssa.benefit_levers.payable.share": {"2024": 0.8}, + "gov.contrib.ssa.benefit_levers.payable.oasi_only": {"2024": True}, + } + ) + simulation = HouseholdMicrosimulation( + situation=situation({f"{public}_base": {multi_year: 6_000}}), + reform=reform, + tax_benefit_system=system, + ) + simulation.set_input("household_weight", multi_year, [1]) + # A positive exact base lacks annual base inputs. The adapter itself must + # reject the unspecified enum, rather than comparing decoded display text. + with pytest.raises(ValueError, match="dependent benefit fund"): + simulation.calculate(public, multi_year) + for fund, expected in (("di", 6_000), ("oasi", 4_800)): + branch = simulation.get_branch(f"multi_year_{fund}") + for year in (2024, 2025): + branch.set_input(f"{public}_fund", year, [fund]) + drop_inherited_values(branch) + assert amount(branch, public, multi_year) == pytest.approx(expected) diff --git a/policyengine_us/tools/social_security_periods.py b/policyengine_us/tools/social_security_periods.py index fbfa06b458e..ba3b60eda6d 100644 --- a/policyengine_us/tools/social_security_periods.py +++ b/policyengine_us/tools/social_security_periods.py @@ -3,6 +3,7 @@ import numpy as np from policyengine_core.periods import YEAR from policyengine_core.periods import period as to_period +from policyengine_core.simulations import Simulation as CoreSimulation from policyengine_us.tools.social_security_benefit_levers import ( SocialSecurityBenefitSource, @@ -54,7 +55,7 @@ def _calculate(self, variable_name, period=None): base_abolished = False if has_exact_base or base_variable.is_neutralized or base_abolished: # Let Core preserve exact base-input and neutralization precedence. - amount = self.calculate(base_name, period) + amount = CoreSimulation.calculate(self, base_name, period) else: if any( known.start <= period.stop and known.stop >= period.start @@ -64,7 +65,7 @@ def _calculate(self, variable_name, period=None): "A multi-year reported input with base overrides for other " "periods requires separately supplied annual source amounts" ) - amount = self.calculate(reported_name, period) + amount = CoreSimulation.calculate(self, reported_name, period) nonzero = amount != 0 # Core accepts a stored multi-year YEAR input before checking period # consistency. A uniform multiplier needs no annual allocation, while @@ -72,7 +73,7 @@ def _calculate(self, variable_name, period=None): common_factor = None for year_number in range(period.start.year, period.stop.year + 1): year = to_period(year_number) - source = self.calculate(f"{variable_name}_source", year) + source = CoreSimulation.calculate(self, f"{variable_name}_source", year) if np.any(source != SocialSecurityBenefitSource.reported): raise NotImplementedError( "Multi-year Social Security sources require the reported route; " @@ -85,7 +86,7 @@ def _calculate(self, variable_name, period=None): "social_security_across_the_board_factor", f"{variable_name}_payable_factor", ): - factor = factor * self.calculate(name, year) + factor = factor * CoreSimulation.calculate(self, name, year) try: annual_base_abolished = self.tax_benefit_system.parameters( year @@ -99,9 +100,16 @@ def _calculate(self, variable_name, period=None): and self.tax_benefit_system.parameters( year ).gov.contrib.ssa.benefit_levers.payable.oasi_only - and np.any(self.calculate("social_security_payable_factor", year) != 1) + and np.any( + CoreSimulation.calculate( + self, "social_security_payable_factor", year + ) + != 1 + ) ): - fund = self.calculate("social_security_dependents_fund", year) + fund = CoreSimulation.calculate( + self, "social_security_dependents_fund", year + ) if np.any((amount > 0) & (fund == SocialSecurityFund.unspecified)): raise ValueError( "OASI-only payable cuts require the dependent benefit fund" From 849ad59d806701efc98ce6988be2349f951f2ca0 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 16:24:43 -0400 Subject: [PATCH 17/76] Preserve legacy Social Security inputs across ingress and extraction --- changelog.d/ss-benefit-levers.added.md | 2 +- docs/social-security-benefit-levers.md | 23 ++- policyengine_us/system.py | 8 +- .../test_social_security_inputs.py | 150 ++++++++++++++++-- .../tools/social_security_inputs.py | 133 +++++++++++++++- 5 files changed, 290 insertions(+), 26 deletions(-) diff --git a/changelog.d/ss-benefit-levers.added.md b/changelog.d/ss-benefit-levers.added.md index 3720cded084..d7ecf8a8dbe 100644 --- a/changelog.d/ss-benefit-levers.added.md +++ b/changelog.d/ss-benefit-levers.added.md @@ -1 +1 @@ -Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. +Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs migrate on every ingress route; later public overrides seed missing sources and survive extraction. diff --git a/docs/social-security-benefit-levers.md b/docs/social-security-benefit-levers.md index 9a7e974e9d6..d825af71c83 100644 --- a/docs/social-security-benefit-levers.md +++ b/docs/social-security-benefit-levers.md @@ -19,13 +19,22 @@ inputs, including zero, with the original effective per-capita SOI/CBO uprating. Public formulas have no variable-level uprating. Dataset extension still uprates original column names before constructor migration. -Both PE-US constructors migrate a legacy public input to its reported twin -separately for each supplied period. When both names are supplied in the same -period, the reported amount remains the source and the public input remains a -final override. Later public `set_input` calls and aggregate `social_security` -inputs are final overrides. Core `SimulationBuilder` and YAML bypass constructor -migration: use `_reported` for lever-sensitive fixtures. A public override in -one year does not become the reported source for later years. +Initial ingress through either PE-US constructor, dataset loading, or Core's +`SimulationBuilder` (including YAML) migrates a legacy public input to its +reported twin separately for each supplied period. The public formula then +applies the levers in the stored year as well as later years. When both names +are supplied in the same period, the reported amount remains the source and +the public input remains a final override. + +A public `set_input` after construction is a final override for its own period. +It also seeds the reported twin for that period unless an actual twin input +already exists; a calculated or uprated cache does not count as an input. +Later years uprate that source. An existing reported input, including zero, +retains precedence. Aggregate `social_security` inputs remain final overrides. +Clones and branches retain these rules. Later public writes replace their +runtime input-name list privately, preserving overrides in legacy extraction +without changing a parent or sibling clone. Migrated inputs extract as twins; +later overrides extract alongside their twins and retain precedence on reload. Each component has a `_source` selector with default `reported`. Reserved options `supplied_aime`, `supplied_pia`, `supplied_ratio_survivor`, and `history` diff --git a/policyengine_us/system.py b/policyengine_us/system.py index 22c3d707221..6e82e2e0635 100644 --- a/policyengine_us/system.py +++ b/policyengine_us/system.py @@ -36,7 +36,10 @@ ) from .tools.default_uprating import add_default_uprating from .tools.period_branch import delete_input_values -from .tools.social_security_inputs import migrate_social_security_inputs +from .tools.social_security_inputs import ( + install_social_security_input_compatibility, + migrate_social_security_inputs, +) from .tools.social_security_periods import SocialSecurityPeriodMixin from .tools.section_911 import validate_section_911_batch_inputs from .tools.per_capita_uprating import ( @@ -57,6 +60,8 @@ create_spm_provider, ) +install_social_security_input_compatibility() + COUNTRY_DIR = Path(__file__).parent CURRENT_YEAR = 2024 @@ -97,6 +102,7 @@ class CountryTaxBenefitSystem(TaxBenefitSystem): """ variables_dir = COUNTRY_DIR / "variables" + _social_security_input_compatibility = True auto_carry_over_input_variables = True basic_inputs = [ "state_name", diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index 4059a7c773c..49fe6da84b4 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -1,8 +1,8 @@ """Input compatibility, branch isolation and complete source round trips. -These cases need Python because YAML bypasses PE-US constructors and cannot -exercise cloning, stored-input bookkeeping or source export. They reuse the -read-only country model and tiny populations, without loading external data. +These cases cover Core's YAML builder as well as cloning, stored-input +bookkeeping and source export. They reuse the read-only country model and +tiny populations, without loading external data. """ import json @@ -52,6 +52,25 @@ def make_simulation(wrapper=Simulation, inputs=None): return wrapper(situation=situation(inputs or {}), tax_benefit_system=system) +def make_core_simulation(inputs=None, route="dict_entities"): + builder = SimulationBuilder() + inputs = inputs or {} + if route == "dict_entities": + return builder.build_from_dict(system, situation(inputs)) + if route == "entities": + return builder.build_from_entities(system, situation(inputs)) + variables = { + "age": {2024: [70]}, + **{ + name: {period: [value] for period, value in values.items()} + for name, values in inputs.items() + }, + } + if route == "dict_variables": + return builder.build_from_dict(system, variables) + return builder.build_from_variables(system, variables) + + def amount(simulation, name, period): return np.asarray(simulation.calculate(name, period))[0] @@ -120,23 +139,110 @@ def test_later_public_override_survives_clones_and_cache_clear(wrapper): assert amount(simulation, f"{public}_reported", 2024) == 12_000 -def test_core_builder_public_inputs_are_final_overrides(): +@pytest.mark.parametrize( + "route", ("dict_entities", "entities", "dict_variables", "variables") +) +def test_core_builder_migrates_initial_public_inputs_and_preserves_precedence(route): retirement, disability = SOCIAL_SECURITY_COMPONENTS[:2] - simulation = SimulationBuilder().build_from_dict( - system, - situation( - { - f"{retirement}_reported": {2024: 12_000}, - disability: {2024: 777}, - FACTOR: {2024: 1.1}, - } - ), + simulation = make_core_simulation( + { + retirement: {2024: 12_000}, + disability: {2024: 777}, + f"{disability}_reported": {2024: 666}, + FACTOR: {2024: 1.1, 2026: 1.08}, + }, + route, ) assert amount(simulation, retirement, 2024) == pytest.approx(13_200) + assert input_periods(simulation, retirement) == set() + assert input_periods(simulation, f"{retirement}_reported") == {"2024"} + assert retirement not in simulation.input_variables + assert f"{retirement}_reported" in simulation.input_variables + later_source = amount(simulation, f"{retirement}_reported", 2026) + assert later_source == pytest.approx(13_522.16, abs=0.01) + assert amount(simulation, retirement, 2026) == pytest.approx(later_source * 1.08) + # Existing both-name inputs retain the established explicit-override rule. assert amount(simulation, disability, 2024) == 777 + assert amount(simulation, f"{disability}_reported", 2024) == 666 assert input_periods(simulation, disability) == {"2024"} + assert disability in simulation.input_variables + assert "age" in simulation.input_variables + payload = simulation.extract_person(0) + person = next(iter(payload["people"].values())) + assert person[disability] == {"2024": 777} + assert person[f"{disability}_reported"] == {"2024": 666} + assert "age" in person simulation.set_input(retirement, 2024, [0]) assert amount(simulation, retirement, 2024) == 0 + assert amount(simulation, f"{retirement}_reported", 2024) == 12_000 + + +@pytest.mark.parametrize("wrapper", (*WRAPPERS, make_core_simulation)) +@pytest.mark.parametrize("reported_input", (None, 0, 6_000)) +def test_later_public_write_seeds_only_absent_reported_input(wrapper, reported_input): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + inputs = {} if reported_input is None else {reported: {2024: reported_input}} + simulation = ( + wrapper(inputs) + if wrapper is make_core_simulation + else make_simulation(wrapper, inputs) + ) + # A calculated twin, even a cached zero, is not an explicitly supplied twin. + assert amount(simulation, reported, 2024) == (reported_input or 0) + simulation.set_input(FACTOR, 2024, [1.08]) + simulation.set_input(public, 2024, [12_000]) + expected_source = 12_000 if reported_input is None else reported_input + assert amount(simulation, public, 2024) == 12_000 + assert amount(simulation, reported, 2024) == expected_source + assert input_periods(simulation, reported) == {"2024"} + assert public in simulation.input_variables + assert reported in simulation.input_variables + assert amount(simulation, public, 2026) == pytest.approx( + expected_source * 13_522.16 / 12_000, abs=0.01 + ) + + +@pytest.mark.parametrize("wrapper", (*WRAPPERS, make_core_simulation)) +def test_later_monthly_public_write_seeds_the_stored_rolling_year(wrapper): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + simulation = ( + wrapper() if wrapper is make_core_simulation else make_simulation(wrapper) + ) + simulation.set_input(public, "2024-07", [12_000]) + # Core retains a rolling annual input for a monthly annual-flow write. + assert input_periods(simulation, public) == {"year:2024-07"} + assert input_periods(simulation, reported) == {"year:2024-07"} + assert simulation.get_array(public, "year:2024-07")[0] == 12_000 + assert simulation.get_array(reported, "year:2024-07")[0] == 12_000 + assert amount(simulation, public, "2024-07") == 1_000 + assert amount(simulation, public, 2026) > 12_000 + + +@pytest.mark.parametrize("wrapper", WRAPPERS) +def test_clone_and_branch_later_public_writes_keep_input_lists_private(wrapper): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + simulation = make_simulation(wrapper, {public: {2024: 12_000}}) + original_variables = list(simulation.input_variables) + original_keys = set(simulation._user_input_keys) + parent_later_source = amount(simulation, reported, 2025) + for copy in ( + simulation.clone(clone_tax_benefit_system=False), + simulation.get_branch("new_public_override"), + ): + copy.set_input(public, 2025, [24_000]) + drop_inherited_values(copy) + assert amount(copy, public, 2025) == 24_000 + assert amount(copy, reported, 2025) == 24_000 + assert amount(copy, public, 2026) > 24_000 + assert public in copy.input_variables + assert copy.input_variables is not simulation.input_variables + assert simulation.input_variables == original_variables + assert simulation._user_input_keys == original_keys + assert amount(simulation, public, 2025) == parent_later_source + assert input_periods(simulation, reported) == {"2024"} def test_dataset_constructor_migrates_stored_original_columns(): @@ -202,7 +308,7 @@ def test_complete_snapshot_preserves_histories_and_separate_final_overrides(): simulation.set_input(public, 2027, [0]) payload = json.loads(json.dumps(export_source_inputs(simulation))) person = payload["sources"]["people"]["person"] - assert set(person[f"{public}_reported"]) == {"2024", "2026"} + assert set(person[f"{public}_reported"]) == {"2024", "2026", "2027"} assert public not in person assert set(payload["overrides"][public]) == {"2026", "2027"} assert person["ss_aime_input"] == {"2024": 1_000, "2026": 2_000} @@ -232,6 +338,22 @@ def test_legacy_extraction_uses_reported_sources_and_months_keep_flow_conversion assert amount(copy, public, "2024-07") == 1_000 +@pytest.mark.parametrize("wrapper", WRAPPERS) +@pytest.mark.parametrize("override", (0, 24_000)) +def test_legacy_extraction_preserves_later_public_override(wrapper, override): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + simulation = make_simulation(wrapper, {public: {2024: 12_000}}) + simulation.set_input(public, 2024, [override]) + extracted = simulation.extract_person(0) + person = next(iter(extracted["people"].values())) + assert person[public]["2024"] == override + assert person[reported]["2024"] == 12_000 + reloaded = Simulation(situation=extracted, tax_benefit_system=system) + assert amount(reloaded, public, 2024) == override + assert amount(reloaded, reported, 2024) == 12_000 + + def test_reported_twins_are_pure_and_original_dataset_names_keep_their_uprater(): for public in SOCIAL_SECURITY_COMPONENTS: original = system.get_variable(public) diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index b45cc864e98..3f94ce1c977 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -1,10 +1,12 @@ """Compatibility ingress and complete, source-safe simulation snapshots. -Constructor migration is intentionally separate from ``set_input``: a later -write to a public benefit is an output override, including an explicit zero. +Initial ingress moves legacy benefit inputs to their pure reported sources. +A later public write overrides its period and seeds an absent reported source. """ import json +from contextvars import ContextVar +from functools import wraps from policyengine_core.simulations.simulation import NpEncoder @@ -19,6 +21,131 @@ ) +_initial_social_security_ingress = ContextVar( + "initial_social_security_ingress", default=False +) + + +def install_social_security_input_compatibility(): + """Extend Core's ingress routes only for systems opting into this policy. + + Core's entity builder writes holders directly, and its variable-level + ``set_input`` handler runs only for mismatched period units. A builder + adapter therefore migrates initial sources, while a simulation adapter + preserves subsequent public overrides, including on raw Core simulations. + """ + from policyengine_core.simulations import Simulation as CoreSimulation + from policyengine_core.simulations.simulation_builder import SimulationBuilder + + original_set_input = CoreSimulation.set_input + if not getattr(original_set_input, "_social_security_compatibility", False): + + @wraps(original_set_input) + def set_input(simulation, variable_name, period, value): + track_public_write = ( + variable_name in SOCIAL_SECURITY_COMPONENTS + and getattr( + simulation.tax_benefit_system, + "_social_security_input_compatibility", + False, + ) + and hasattr(simulation, "input_variables") + and not _initial_social_security_ingress.get() + ) + if not track_public_write: + return original_set_input(simulation, variable_name, period, value) + + holder = simulation.get_holder(variable_name) + branch = simulation.branch_name + before = { + supplied_period: holder._stored_sequence_number(supplied_period, branch) + for supplied_period in simulation._get_set_input_periods(variable_name) + } + result = original_set_input(simulation, variable_name, period, value) + reported_name = f"{variable_name}_reported" + reported_periods = set(simulation._get_set_input_periods(reported_name)) + # Period helpers can store a different period from the one passed + # in: a monthly input to a yearly variable stores a rolling year. + # Inspect actual stores rather than guessing those period aliases. + for supplied_period in simulation._get_set_input_periods(variable_name): + if ( + holder._stored_sequence_number(supplied_period, branch) + == before.get(supplied_period) + or supplied_period in reported_periods + ): + continue + supplied_value = simulation._get_recorded_input_array( + variable_name, supplied_period + ) + if supplied_value is not None: + original_set_input( + simulation, + reported_name, + supplied_period, + supplied_value.copy(), + ) + + # Core clones share this legacy list. Replace it privately so a + # later override is extracted without changing the source clone. + simulation.input_variables = list( + dict.fromkeys( + [*simulation.input_variables, variable_name, reported_name] + ) + ) + return result + + set_input._social_security_compatibility = True + CoreSimulation.set_input = set_input + + def wrap_builder(original_build): + @wraps(original_build) + def build(builder, tax_benefit_system, input_dict, simulation=None): + if not getattr( + tax_benefit_system, "_social_security_input_compatibility", False + ): + return original_build( + builder, tax_benefit_system, input_dict, simulation + ) + token = _initial_social_security_ingress.set(True) + try: + result = original_build( + builder, tax_benefit_system, input_dict, simulation + ) + finally: + _initial_social_security_ingress.reset(token) + # Only build_from_dict rebuilds Core's legacy extraction list. + # Direct entity/variable builders need the same recorded inputs, + # including non-benefit sources and exact both-name overrides. + visible_branches = result._get_visible_branch_names() + recorded_names = { + name + for name, branch_name, _ in result._user_input_keys + if branch_name in visible_branches + } + result.input_variables = list( + dict.fromkeys( + [ + *getattr(result, "input_variables", []), + *sorted( + name + for name in recorded_names + if result._get_set_input_periods(name) + ), + ] + ) + ) + migrate_social_security_inputs(result) + return result + + build._social_security_compatibility = True + return build + + for method_name in ("build_from_entities", "build_from_variables"): + original_build = getattr(SimulationBuilder, method_name) + if not getattr(original_build, "_social_security_compatibility", False): + setattr(SimulationBuilder, method_name, wrap_builder(original_build)) + + def migrate_social_security_inputs(simulation, *, include_baseline=True): """Move legacy component inputs, period by period, into pure sources. @@ -77,7 +204,7 @@ def migrate_social_security_inputs(simulation, *, include_baseline=True): # Core's legacy extraction reads this runtime list, which was assembled # before migration. Retain public names only where actual overrides remain. runtime_inputs = [] - for name in simulation.input_variables: + for name in getattr(simulation, "input_variables", []): if name not in migrated or simulation._get_set_input_periods(name): runtime_inputs.append(name) if name in migrated: From 1a7b9124a3e9278081fae0bf04064f43c860b85b Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 16:28:01 -0400 Subject: [PATCH 18/76] Compare legacy benefit ingress against unmodified main --- .../test_social_security_current_law.py | 731 ++++++++++-------- 1 file changed, 388 insertions(+), 343 deletions(-) diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index c015b5e596e..6772d7cd531 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -1,47 +1,37 @@ -"""Differential compatibility with the four original public benefit inputs. - -Python is required for constructor, Core builder, dataset, cache/branch and -export APIs that YAML cannot exercise. The reference restores the original -four pure input definitions and disables only the new constructor migration; -all other current-law policy and wrapper behavior is identical. One model -clone is shared read-only across the module. Every population has two people, -with no external dataset, and only the neutral-reform branches copy policy. -Local runtime/memory measurements are reported with the PR's test receipt. +"""Compare legacy public ingress with an unmodified main checkout. + +The same probe runs in separate Python processes against this checkout and +main. No variable definitions or input names are substituted in the reference. +Set PE_US_MAIN_CHECKOUT to a clean main checkout, or make a local main remote +ref available (the fixture archives it without changing either checkout). """ -from copy import deepcopy -from unittest.mock import patch +import io +import json +import os +import subprocess +import sys +import tarfile +import tempfile +from pathlib import Path import numpy as np -import pandas as pd import pytest -from policyengine_core.periods import YEAR -from policyengine_core.reforms import Reform -from policyengine_core.simulations import SimulationBuilder -from policyengine_core.variables import Variable - -from policyengine_us import Microsimulation, Simulation -from policyengine_us.system import system -from policyengine_us.tools.period_branch import drop_inherited_values -from policyengine_us.tools.social_security_inputs import ( - SOCIAL_SECURITY_COMPONENTS, - export_source_inputs, - reload_source_inputs, -) - -# All six social_security* variables present before this change, including -# the payroll-tax base whose name shares the benefit prefix. -PREEXISTING_SOCIAL_SECURITY = ( - *SOCIAL_SECURITY_COMPONENTS, - "social_security", - "social_security_taxable_self_employment_income", +COMPONENTS = ( + "social_security_retirement", + "social_security_disability", + "social_security_survivors", + "social_security_dependents", ) +BENEFITS = (*COMPONENTS, "social_security") DOWNSTREAM = ( + "social_security_taxable_self_employment_income", "taxable_social_security", "tax_unit_taxable_social_security", "adjusted_gross_income", "income_tax", + "household_net_income", "employee_social_security_tax", "employer_social_security_tax", "self_employment_social_security_tax", @@ -54,46 +44,15 @@ "medicare_irmaa_magi", "medicare_part_b_premium", ) -OUTPUTS = (*PREEXISTING_SOCIAL_SECURITY, *DOWNSTREAM) -NAMES = ("retiree", "disabled") -AMOUNTS = ((24_000, 0, 1_800, 600), (0, 6_000, 1_200, 300)) - - -class HouseholdMicrosimulation(Microsimulation): - default_dataset = None - - -@pytest.fixture(scope="module") -def original_input_system(): - """Restore main's input contracts, including effective per-capita growth.""" - original = system.clone() - for public_name in SOCIAL_SECURITY_COMPONENTS: - reported = original.variables[f"{public_name}_reported"] - input_class = type( - public_name, - (Variable,), - { - "value_type": reported.value_type, - "entity": reported.entity, - "definition_period": YEAR, - "unit": reported.unit, - "label": public_name, - "uprating": reported.uprating, - }, - ) - original.replace_variable(input_class) - restored = original.variables[public_name] - assert restored.is_input_variable() - assert restored.uprating == reported.uprating - assert restored.uprating.endswith("social_security_per_capita") - return original +AMOUNTS = ((12_000, 0, 1_800, 600), (0, 6_000, 1_200, 300)) -def situation(*, source="public", years=(2024, 2026), include_benefits=True): - people = {} - for index, name in enumerate(NAMES): +def situation(*, include_benefits=True): + """Identical legacy public-name inputs for every checkout and wrapper.""" + result = {"people": {}} + for index, name in enumerate(("retiree", "disabled")): person = { - "age": {year: (70 if index == 0 else 45) + year - 2024 for year in years}, + "age": {2024: 70 if index == 0 else 45, 2026: 72 if index == 0 else 47}, "employment_income": {2024: 15_000 if index == 0 else 0}, "self_employment_income": {2024: 3_000 if index == 0 else 0}, "taxable_interest_income": {2024: 50_000 if index == 0 else 10}, @@ -101,294 +60,380 @@ def situation(*, source="public", years=(2024, 2026), include_benefits=True): "takes_up_ssi_if_eligible": {2024: True}, } if include_benefits: - for public_name, amount in zip(SOCIAL_SECURITY_COMPONENTS, AMOUNTS[index]): - name_to_store = ( - f"{public_name}_reported" if source == "reported" else public_name - ) - person[name_to_store] = { - year: amount * (1 if year == 2024 else 1.07) for year in years - } - people[name] = person - result = {"people": people} + for component, amount in zip(COMPONENTS, AMOUNTS[index]): + person[component] = {2024: amount} + result["people"][name] = person for plural in ("households", "tax_units", "spm_units", "families", "marital_units"): - result[plural] = {name: {"members": [name]} for name in NAMES} + result[plural] = {name: {"members": [name]} for name in result["people"]} for household in result["households"].values(): - household["state_code"] = {year: "TX" for year in years} + household["state_code"] = {2024: "TX", 2026: "TX"} for spm_unit in result["spm_units"].values(): spm_unit["takes_up_snap_if_eligible"] = {2024: True} return result -def reference_simulation(original_input_system, wrapper=Simulation, **kwargs): - # Both wrappers retain all their existing income/geography preparation. - # Only this PR's compatibility migration is absent from the reference. - with patch( - "policyengine_us.system.migrate_social_security_inputs", return_value=set() - ): - return wrapper(tax_benefit_system=original_input_system, **kwargs) +def _snapshot(simulation, periods, variables=BENEFITS): + return { + f"{variable}@{period}": np.asarray( + simulation.calculate(variable, period) + ).tolist() + for period in periods + for variable in variables + } -def assert_equal_outputs(current, original, requests, variables=OUTPUTS): - for period in requests: - for variable in variables: - np.testing.assert_array_equal( - np.asarray(current.calculate(variable, period)), - np.asarray(original.calculate(variable, period)), - err_msg=f"Current-law difference for {variable}, {period}", - ) - +def _write_h5(path): + """Variable-centric v2pop-style H5, with public columns in its scored year.""" + import h5py -@pytest.mark.parametrize( - "requests", - [ - ("2024-12", 2026, "2025-03", 2024, 2025, "2026-08"), - ("2026-08", 2025, 2024, "2025-03", 2026, "2024-12"), - ], - ids=("month-first-forward", "month-first-reverse"), -) -def test_household_current_law_matches_original_through_years_and_months( - original_input_system, requests -): - inputs = situation() - original = reference_simulation(original_input_system, situation=deepcopy(inputs)) - current = Simulation(situation=inputs, tax_benefit_system=system) - assert_equal_outputs(current, original, requests) - - # The supplied 2024 source must be read again for 2026 IRMAA, even after - # calculating 2026 income and a different source amount for that year. - np.testing.assert_array_equal( - current.calculate("medicare_irmaa_magi_two_years_prior", 2026), - current.calculate("medicare_irmaa_tax_year_magi", 2024), - ) - assert not np.array_equal( - current.calculate("medicare_irmaa_tax_year_magi", 2026), - current.calculate("medicare_irmaa_tax_year_magi", 2024), - ) - - -def dataset_frame(): - """The dataset loader itself sets the stored public columns as inputs.""" - data = {"person_id": [1, 2]} + ids = np.array([1, 2]) + columns = {"person_id": ids, "age": [72, 47]} for entity in ("household", "tax_unit", "spm_unit", "family", "marital_unit"): - data[f"person_{entity}_id"] = [1, 2] - data[f"{entity}_id"] = [1, 2] - for year in (2024, 2026): - data[f"age__{year}"] = [70 + year - 2024, 45 + year - 2024] - data[f"state_fips__{year}"] = [48, 48] - for component_index, public_name in enumerate(SOCIAL_SECURITY_COMPONENTS): - data[f"{public_name}__{year}"] = [ - amounts[component_index] * (1 if year == 2024 else 1.07) - for amounts in AMOUNTS - ] - data["employment_income__2024"] = [15_000, 0] - data["self_employment_income__2024"] = [3_000, 0] - data["taxable_interest_income__2024"] = [50_000, 10] - data["is_disabled__2024"] = [False, True] - data["takes_up_ssi_if_eligible__2024"] = [True, True] - data["takes_up_snap_if_eligible__2024"] = [True, True] - return pd.DataFrame(data) - - -def test_dataset_loader_current_law_matches_original(original_input_system): - frame = dataset_frame() - original = reference_simulation( - original_input_system, Microsimulation, dataset=frame.copy(deep=True) - ) - current = Microsimulation(dataset=frame, tax_benefit_system=system) - assert_equal_outputs(current, original, ("2025-06", 2026, 2024, 2025)) - - -@pytest.mark.parametrize("ingress", ["core_builder", "set_input"]) -def test_core_builder_and_dataset_style_source_writes_match_original( - original_input_system, ingress -): - if ingress == "core_builder": - original_situation = situation() - current_situation = situation(source="reported") - else: - original_situation = situation(include_benefits=False) - current_situation = deepcopy(original_situation) - original = SimulationBuilder().build_from_dict( - original_input_system, original_situation + columns[f"person_{entity}_id"] = ids + columns[f"{entity}_id"] = ids + columns[f"{entity}_weight"] = [1.0, 1.0] + columns["state_fips"] = [48, 48] + for index, component in enumerate(COMPONENTS): + columns[component] = [amounts[index] for amounts in AMOUNTS] + with h5py.File(path, "w") as stored: + for name, values in columns.items(): + stored.create_dataset(f"{name}/2026", data=values) + + +def _probe(): + """Only imports from the checkout selected through the subprocess path.""" + import gc + + import yaml + from policyengine_core.data import Dataset + from policyengine_core.reforms import Reform + from policyengine_core.simulations import SimulationBuilder + from policyengine_us import Microsimulation, Simulation + from policyengine_us.system import system + + class HouseholdMicrosimulation(Microsimulation): + default_dataset = None + + has_levers = "social_security_retirement_reported" in system.variables + neutral = ( + Reform.from_dict( + {"gov.contrib.ssa.benefit_levers.across_the_board.rate": {"2024": 0}} + ) + if has_levers + else None ) - current = SimulationBuilder().build_from_dict(system, current_situation) - if ingress == "set_input": - for component_index, public_name in enumerate(SOCIAL_SECURITY_COMPONENTS): - for year in (2024, 2026): - values = [ - amounts[component_index] * (1 if year == 2024 else 1.07) - for amounts in AMOUNTS - ] - original.set_input(public_name, year, values) - current.set_input(f"{public_name}_reported", year, values) - assert_equal_outputs(current, original, ("2026-11", 2025, 2024, 2026)) - - -def test_clones_and_neutral_reform_branches_match_original(original_input_system): - original = reference_simulation(original_input_system, situation=situation()) - current = Simulation(situation=situation(), tax_benefit_system=system) - assert_equal_outputs(current, original, (2024, 2026)) - - # Warm-cache clones and ordinary branches should preserve the source - # uprating path as well as any final public override applied later. - for current_copy, original_copy in ( - ( - current.clone(clone_tax_benefit_system=False), - original.clone(clone_tax_benefit_system=False), - ), - (current.get_branch("copy"), original.get_branch("copy")), + result = {} + for label, wrapper in ( + ("simulation", Simulation), + ("microsimulation", HouseholdMicrosimulation), ): - current_copy.set_input("social_security_retirement", 2026, [0, 123]) - original_copy.set_input("social_security_retirement", 2026, [0, 123]) - drop_inherited_values(current_copy) - drop_inherited_values(original_copy) - assert_equal_outputs(current_copy, original_copy, (2026, "2025-07")) - - # Register an actual no-change parameter reform after outputs were - # cached. This exercises reform replay and invalidation, while copying - # policy only for these two branches and preserving the read-only models. - neutral = Reform.from_dict( - {"gov.contrib.ssa.benefit_levers.across_the_board.rate": {"2024": 0}} - ) - original_reform = original.get_branch("neutral_reform", clone_system=True) - current_reform = current.get_branch("neutral_reform", clone_system=True) - original_reform.apply_reform(neutral) - current_reform.apply_reform(neutral) - assert_equal_outputs(current_reform, original_reform, (2025, 2026)) - assert_equal_outputs(current, original, (2024, 2026)) - - -def test_legacy_extraction_reload_matches_original(original_input_system): - inputs = situation(years=(2024,)) - original = reference_simulation( - original_input_system, HouseholdMicrosimulation, situation=deepcopy(inputs) - ) - current = HouseholdMicrosimulation(situation=inputs, tax_benefit_system=system) - original_payload = original.extract_person(0) - current_payload = current.extract_person(0) - original_reload = reference_simulation( - original_input_system, situation=original_payload + simulation = wrapper(situation=situation(), tax_benefit_system=system) + result[f"constructor_{label}"] = _snapshot(simulation, (2024, 2026, "2025-07")) + if label == "simulation": + result["constructor_downstream"] = _snapshot( + simulation, (2024, 2026), DOWNSTREAM + ) + assert ( + result["constructor_downstream"][ + "medicare_irmaa_magi_two_years_prior@2026" + ] + == result["constructor_downstream"]["medicare_irmaa_tax_year_magi@2024"] + ) + del simulation + gc.collect() + + # A fresh period has no explicit twin. The later public override + # must be final in 2026 and seed main's uprating path into 2027. + simulation = wrapper(situation=situation(), tax_benefit_system=system) + simulation.set_input("social_security_retirement", 2026, [24_000, 0]) + result[f"later_override_{label}"] = _snapshot(simulation, (2026, 2027)) + if label == "simulation": + # The complete-history helper is new in this change. On main, + # the equivalent reference is the same live source/override state. + try: + from policyengine_us.tools.social_security_inputs import ( + export_source_inputs, + reload_source_inputs, + ) + except ImportError: + result["complete_export_reload"] = _snapshot( + simulation, (2024, 2025, 2026, 2027) + ) + else: + payload = export_source_inputs(simulation) + assert payload["overrides"]["social_security_retirement"]["2026"] == [ + 24_000, + 0, + ] + reloaded = reload_source_inputs(payload, tax_benefit_system=system) + result["complete_export_reload"] = _snapshot( + reloaded, (2024, 2025, 2026, 2027) + ) + assert result["complete_export_reload"] == _snapshot( + simulation, (2024, 2025, 2026, 2027) + ) + del reloaded + del simulation + gc.collect() + + # Same-period explicit twin precedence remains covered by the input + # tests. This regression specifically checks the later final write + # survives legacy extraction, including its tax and net-income effects. + simulation = wrapper(situation=situation(), tax_benefit_system=system) + simulation.set_input("social_security_retirement", 2024, [24_000, 0]) + # Extract before calculations introduce earlier lookback cache periods + # into Core's legacy first-known-period selection. + payload = simulation.extract_person(0) + reloaded = Simulation(situation=payload, tax_benefit_system=system) + live = _snapshot(simulation, (2024,), (*BENEFITS, *DOWNSTREAM)) + extracted = _snapshot(reloaded, (2024,), (*BENEFITS, *DOWNSTREAM)) + assert extracted == {key: values[:1] for key, values in live.items()} + result[f"extraction_reload_{label}"] = extracted + del simulation, reloaded + gc.collect() + + # Core constructs a baseline before the US wrapper returns. The + # current tree's neutral lever reform must migrate that clone too; + # unmodified main supplies the same current-law values directly. + inputs = situation() + for person in inputs["people"].values(): + # Existing wage/self-employment constructor migrations run after + # Core creates its baseline. Isolate the benefit ingress here. + person.pop("employment_income") + person.pop("self_employment_income") + simulation = wrapper( + situation=inputs, tax_benefit_system=system, reform=neutral + ) + baseline = simulation.baseline if has_levers else simulation + assert baseline is not None + result[f"automatic_baseline_{label}"] = _snapshot( + baseline, (2024, 2026, "2025-07") + ) + result[f"neutral_constructor_{label}"] = _snapshot(simulation, (2024, 2026)) + del baseline, simulation + gc.collect() + + # YAML reaches Core's builder directly, bypassing both US constructors. + yaml_inputs = yaml.safe_load(yaml.safe_dump(situation())) + simulation = SimulationBuilder().build_from_dict(system, yaml_inputs) + result["yaml_core_builder"] = _snapshot(simulation, (2024, 2026)) + result["yaml_core_downstream"] = _snapshot(simulation, (2026,), DOWNSTREAM) + del simulation + gc.collect() + + simulation = SimulationBuilder().build_from_dict( + system, situation(include_benefits=False) ) - current_reload = Simulation(situation=current_payload, tax_benefit_system=system) - assert_equal_outputs(current_reload, original_reload, ("2024-10", 2024, 2025)) - for public_name, amount in zip(SOCIAL_SECURITY_COMPONENTS, AMOUNTS[0]): - assert current_reload.calculate(public_name, 2024)[0] == amount - - -def test_complete_source_and_final_override_reload_matches_original( - original_input_system, -): - original = reference_simulation(original_input_system, situation=situation()) - current = Simulation(situation=situation(), tax_benefit_system=system) - # An unsupplied, calculated year must remain a calculation after export. - assert_equal_outputs(current, original, (2025,)) - for simulation in (original, current): - simulation.set_input("social_security_retirement", 2027, [0, 123]) - simulation.set_input("social_security", 2027, [9_999, 0]) - original_payload = export_source_inputs(original) - current_payload = export_source_inputs(current) - for public_name in SOCIAL_SECURITY_COMPONENTS: - reported = current_payload["sources"]["people"]["retiree"][ - f"{public_name}_reported" - ] - assert set(reported) == {"2024", "2026"} - with patch( - "policyengine_us.system.migrate_social_security_inputs", return_value=set() + for index, component in enumerate(COMPONENTS): + simulation.set_input(component, 2024, [amounts[index] for amounts in AMOUNTS]) + result["public_set_input"] = _snapshot(simulation, (2024, 2026)) + del simulation + gc.collect() + + for route in ("clone", "branch"): + simulation = Simulation(situation=situation(), tax_benefit_system=system) + original_inputs = list(simulation.input_variables) + parent = _snapshot(simulation, (2024, 2026)) + copied = ( + simulation.clone(clone_tax_benefit_system=False) + if route == "clone" + else simulation.get_branch("public_override") + ) + copied.set_input("social_security_retirement", 2026, [24_000, 0]) + result[route] = _snapshot(copied, (2026, 2027)) + assert _snapshot(simulation, (2024, 2026)) == parent + assert simulation.input_variables == original_inputs + if has_levers: + assert copied.input_variables is not simulation.input_variables + result[f"{route}_parent"] = _snapshot(simulation, (2024, 2026)) + if route == "branch": + reformed = simulation.get_branch("neutral_reform", clone_system=True) + if neutral is not None: + reformed.apply_reform(neutral) + result["neutral_reform_branch"] = _snapshot(reformed, (2024, 2026)) + del reformed + del simulation, copied + gc.collect() + + for label, wrapper in ( + ("simulation", Simulation), + ("microsimulation", HouseholdMicrosimulation), ): - original_reload = reload_source_inputs( - original_payload, tax_benefit_system=original_input_system + inputs = situation(include_benefits=False) + for person in inputs["people"].values(): + person["social_security_retirement"] = {2024: 12_000, "2024-07": 500} + person["social_security_disability"] = {"2024-07": 300} + person["social_security_survivors"] = {"year:2024:2": 2_000} + for household in inputs["households"].values(): + household["household_weight"] = {"year:2024:2": 1} + simulation = wrapper(situation=inputs, tax_benefit_system=system) + result[f"monthly_{label}"] = _snapshot( + simulation, + ("2024-07", "year:2024-07", 2024, "2024-06"), + COMPONENTS[:2], ) - current_reload = reload_source_inputs(current_payload, tax_benefit_system=system) - requests = ("2027-08", 2026, 2024, 2025, 2027) - assert_equal_outputs(current_reload, original_reload, requests) - assert_equal_outputs(current_reload, current, requests) - - -def test_later_reported_write_updates_a_fresh_branch_without_changing_parent( - original_input_system, -): - current = Simulation(situation=situation(), tax_benefit_system=system) - original = reference_simulation(original_input_system, situation=situation()) - assert_equal_outputs(current, original, (2024,)) - parent_keys = set(current._user_input_keys) - current_branch = current.get_branch("updated_reported") - original_branch = original.get_branch("updated_original") - current_branch.set_input("social_security_retirement_reported", 2024, [30000, 2400]) - original_branch.set_input("social_security_retirement", 2024, [30000, 2400]) - drop_inherited_values(current_branch) - drop_inherited_values(original_branch) - assert_equal_outputs(current_branch, original_branch, (2024, 2025)) - assert current._user_input_keys == parent_keys - np.testing.assert_array_equal( - current.calculate("social_security_retirement", 2024), [24000, 0] - ) - reloaded = reload_source_inputs( - export_source_inputs(current_branch), tax_benefit_system=system - ) - assert_equal_outputs(reloaded, current_branch, (2024, 2025)) - - -@pytest.mark.parametrize("wrapper", [Simulation, HouseholdMicrosimulation]) -def test_constructor_created_baseline_keeps_migrated_uprating( - original_input_system, - wrapper, -): - # Core creates its automatic baseline before wrapper migration. A source - # present only in 2024 must still uprate into 2025 on that baseline. - neutral = Reform.from_dict( - {"gov.contrib.ssa.benefit_levers.across_the_board.rate": {"2024": 0}} - ) - inputs = situation(years=(2024,)) - # Isolate this source migration from the existing wage/self-employment - # constructor moves, which Core's pre-wrapper baseline does not perform. - for person in inputs["people"].values(): - person.pop("employment_income") - person.pop("self_employment_income") - original = reference_simulation( - original_input_system, wrapper, situation=deepcopy(inputs) - ) - current = wrapper(situation=inputs, tax_benefit_system=system, reform=neutral) - assert current.baseline is not None - assert_equal_outputs(current.baseline, original, (2024, "2025-06", 2025)) - assert_equal_outputs(current, original, (2024, 2025)) - for public in SOCIAL_SECURITY_COMPONENTS: - assert not current.baseline._get_set_input_periods(public) - assert current.baseline._get_set_input_periods(f"{public}_reported") - - -@pytest.mark.parametrize("wrapper", [Simulation, HouseholdMicrosimulation]) -def test_monthly_input_normalization_and_exact_multi_year_reads_match_original( - original_input_system, wrapper -): - inputs = situation(include_benefits=False) - for person in inputs["people"].values(): - person["social_security_retirement"] = {2024: 12_000, "2024-07": 500} - person["social_security_disability"] = {"2024-07": 300} - person["social_security_survivors"] = {"year:2024:2": 2_000} - for household in inputs["households"].values(): - # Main's weighted API also needs an exact weight for this multi-year - # request; source migration does not change Core's weight semantics. - household["household_weight"] = {"year:2024:2": 1} - original = reference_simulation( - original_input_system, wrapper, situation=deepcopy(inputs) - ) - current = wrapper(situation=inputs, tax_benefit_system=system) - # Core's default FLOW input handler maps July to a rolling YEAR starting - # in July. It does not store an exact MONTH on these annual variables. - periods = current._get_set_input_periods("social_security_retirement_reported") - assert {str(period) for period in periods} == {"2024", "year:2024-07"} - assert all(period.unit == YEAR and period.size == 1 for period in periods) - requests = ("2024-07", "year:2024-07", 2024, "2024-06", "year:2024:2") - for public in SOCIAL_SECURITY_COMPONENTS[:2]: - assert_equal_outputs(current, original, requests[:-1], variables=(public,)) - assert_equal_outputs( - current, original, ("year:2024:2",), variables=("social_security_survivors",) - ) - reloaded = reload_source_inputs( - export_source_inputs(current), tax_benefit_system=system - ) - for public in SOCIAL_SECURITY_COMPONENTS[:2]: - assert_equal_outputs(reloaded, original, requests[:-1], variables=(public,)) - assert_equal_outputs( - reloaded, original, ("year:2024:2",), variables=("social_security_survivors",) + result[f"multi_year_{label}"] = _snapshot( + simulation, ("year:2024:2",), ("social_security_survivors",) + ) + del simulation + gc.collect() + + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "public_components_2026.h5" + _write_h5(path) + dataset = Dataset.from_file(path) + assert dataset.data_format == Dataset.TIME_PERIOD_ARRAYS + simulation = Microsimulation(dataset=dataset, tax_benefit_system=system) + result["h5_scored_year"] = _snapshot( + simulation, (2026,), (*BENEFITS, *DOWNSTREAM) + ) + # Main has no benefit-lever parameter. Its role here is the unchanged + # current-law reference; the lever assertion runs in the current tree. + if has_levers: + reform = Reform.from_dict( + { + "gov.contrib.ssa.benefit_levers.across_the_board.start_year": { + "2026": 2026 + }, + "gov.contrib.ssa.benefit_levers.across_the_board.rate": { + "2026": 0.08 + }, + } + ) + reformed = Microsimulation( + dataset=dataset, tax_benefit_system=system, reform=reform + ) + for component in COMPONENTS: + expected = np.asarray(simulation.calculate(component, 2026)) * 1.08 + np.testing.assert_array_equal( + np.asarray(reformed.calculate(component, 2026)), expected + ) + np.testing.assert_allclose( + np.asarray(reformed.calculate("social_security", 2026)), + np.asarray(simulation.calculate("social_security", 2026)) * 1.08, + rtol=1e-7, + atol=0, + ) + return result + + +def _main_checkout(current, directory): + configured = os.environ.get("PE_US_MAIN_CHECKOUT") + if configured: + checkout = Path(configured).resolve() + assert (checkout / "policyengine_us").is_dir(), configured + if (checkout / ".git").exists(): + status = subprocess.check_output( + ["git", "status", "--porcelain", "--untracked-files=no"], + cwd=checkout, + text=True, + ) + assert not status, "PE_US_MAIN_CHECKOUT must have no tracked changes" + return checkout + for ref in ("refs/remotes/origin/main", "refs/remotes/gh/main", "refs/heads/main"): + exists = subprocess.run( + ["git", "rev-parse", "--verify", ref], + cwd=current, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=False, + ) + if exists.returncode == 0: + archive = subprocess.check_output(["git", "archive", ref], cwd=current) + checkout = directory / "main" + checkout.mkdir() + with tarfile.open(fileobj=io.BytesIO(archive)) as stored: + stored.extractall(checkout, filter="data") + return checkout + pytest.skip( + "Differential needs PE_US_MAIN_CHECKOUT or a local origin/main, gh/main, " + "or main ref containing an unmodified main checkout" ) + + +@pytest.fixture(scope="module") +def differential(tmp_path_factory): + directory = tmp_path_factory.mktemp("social_security_differential") + current = Path(__file__).resolve().parents[3] + reference = _main_checkout(current, directory) + results = [] + for label, checkout in (("current", current), ("main", reference)): + destination = directory / f"{label}.json" + environment = dict(os.environ) + environment["PYTHONPATH"] = str(checkout) + for variable in ( + "OMP_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + "MKL_NUM_THREADS", + "VECLIB_MAXIMUM_THREADS", + "NUMEXPR_NUM_THREADS", + ): + environment[variable] = "4" + environment["OMP_WAIT_POLICY"] = "PASSIVE" + completed = subprocess.run( + [ + sys.executable, + str(Path(__file__).resolve()), + "--probe", + str(destination), + ], + cwd=checkout, + env=environment, + capture_output=True, + text=True, + timeout=420, + check=False, + ) + assert completed.returncode == 0, ( + f"{label} differential probe failed:\n" + f"{completed.stdout[-4000:]}\n{completed.stderr[-8000:]}" + ) + results.append(json.loads(destination.read_text())) + return results + + +@pytest.mark.parametrize( + "route", + ( + "constructor_simulation", + "constructor_microsimulation", + "constructor_downstream", + "later_override_simulation", + "later_override_microsimulation", + "extraction_reload_simulation", + "extraction_reload_microsimulation", + "automatic_baseline_simulation", + "automatic_baseline_microsimulation", + "neutral_constructor_simulation", + "neutral_constructor_microsimulation", + "complete_export_reload", + "yaml_core_builder", + "yaml_core_downstream", + "public_set_input", + "clone", + "branch", + "clone_parent", + "branch_parent", + "neutral_reform_branch", + "monthly_simulation", + "monthly_microsimulation", + "multi_year_simulation", + "multi_year_microsimulation", + "h5_scored_year", + ), +) +def test_legacy_public_ingress_matches_unmodified_main(differential, route): + current, main = differential + assert current[route].keys() == main[route].keys() + for request in main[route]: + np.testing.assert_array_equal( + current[route][request], + main[route][request], + err_msg=f"Legacy public ingress differs on {route}: {request}", + ) + + +if __name__ == "__main__": + if len(sys.argv) != 3 or sys.argv[1] != "--probe": + raise SystemExit("Usage: test_social_security_current_law.py --probe OUTPUT") + Path(sys.argv[2]).write_text(json.dumps(_probe())) From 43a07be9ad85df5a03b6db0218b988118ff900a6 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 17:15:42 -0400 Subject: [PATCH 19/76] Retain automatic benefit source provenance across later writes --- changelog.d/ss-benefit-levers.added.md | 2 +- docs/social-security-benefit-levers.md | 14 +- .../test_social_security_current_law.py | 51 +++- .../test_social_security_inputs.py | 225 +++++++++++++++++- .../tools/social_security_inputs.py | 187 ++++++++++++++- 5 files changed, 450 insertions(+), 29 deletions(-) diff --git a/changelog.d/ss-benefit-levers.added.md b/changelog.d/ss-benefit-levers.added.md index d7ecf8a8dbe..cec2a318738 100644 --- a/changelog.d/ss-benefit-levers.added.md +++ b/changelog.d/ss-benefit-levers.added.md @@ -1 +1 @@ -Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs migrate on every ingress route; later public overrides seed missing sources and survive extraction. +Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs migrate on every ingress route; later public overrides refresh automatic sources, preserve explicit reported inputs, and survive extraction. diff --git a/docs/social-security-benefit-levers.md b/docs/social-security-benefit-levers.md index d825af71c83..9b3fd078421 100644 --- a/docs/social-security-benefit-levers.md +++ b/docs/social-security-benefit-levers.md @@ -27,14 +27,20 @@ are supplied in the same period, the reported amount remains the source and the public input remains a final override. A public `set_input` after construction is a final override for its own period. -It also seeds the reported twin for that period unless an actual twin input -already exists; a calculated or uprated cache does not count as an input. -Later years uprate that source. An existing reported input, including zero, -retains precedence. Aggregate `social_security` inputs remain final overrides. +It also seeds or refreshes the reported twin for that period unless the caller +explicitly supplied that twin. Automatically migrated or seeded twins follow +later public writes, preserving main's uprating and earlier-year lookbacks; +calculated or uprated caches do not block seeding. An explicitly supplied +reported input, including zero, retains precedence. Aggregate `social_security` +inputs remain final overrides. Clones and branches retain these rules. Later public writes replace their runtime input-name list privately, preserving overrides in legacy extraction without changing a parent or sibling clone. Migrated inputs extract as twins; later overrides extract alongside their twins and retain precedence on reload. +A private provenance envelope preserves the distinction between automatic and +explicit sources through legacy extraction and complete snapshots, including +subsequent public writes after reload. Builders consume this metadata before +validating ordinary situation inputs. Each component has a `_source` selector with default `reported`. Reserved options `supplied_aime`, `supplied_pia`, `supplied_ratio_survivor`, and `history` diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 6772d7cd531..0dbd6cc0190 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -159,6 +159,8 @@ class HouseholdMicrosimulation(Microsimulation): result["complete_export_reload"] = _snapshot( simulation, (2024, 2025, 2026, 2027) ) + simulation.set_input("social_security_retirement", 2026, [36_000, 0]) + result["complete_export_rewrite"] = _snapshot(simulation, (2026, 2027)) else: payload = export_source_inputs(simulation) assert payload["overrides"]["social_security_retirement"]["2026"] == [ @@ -172,14 +174,23 @@ class HouseholdMicrosimulation(Microsimulation): assert result["complete_export_reload"] == _snapshot( simulation, (2024, 2025, 2026, 2027) ) + reloaded.set_input("social_security_retirement", 2026, [36_000, 0]) + result["complete_export_rewrite"] = _snapshot(reloaded, (2026, 2027)) del reloaded del simulation gc.collect() - # Same-period explicit twin precedence remains covered by the input - # tests. This regression specifically checks the later final write - # survives legacy extraction, including its tax and net-income effects. - simulation = wrapper(situation=situation(), tax_benefit_system=system) + # A same-period later public write updates the twin originally moved + # from public ingress. An explicitly supplied twin retains priority + # in the separate precedence tests. Extraction must preserve this + # origin as well as the final override for a subsequent public write. + inputs = situation() + for person in inputs["people"].values(): + # Legacy extraction does not round-trip the existing earnings + # constructor migrations. Isolate the reported benefit regression. + person.pop("employment_income") + person.pop("self_employment_income") + simulation = wrapper(situation=inputs, tax_benefit_system=system) simulation.set_input("social_security_retirement", 2024, [24_000, 0]) # Extract before calculations introduce earlier lookback cache periods # into Core's legacy first-known-period selection. @@ -187,8 +198,31 @@ class HouseholdMicrosimulation(Microsimulation): reloaded = Simulation(situation=payload, tax_benefit_system=system) live = _snapshot(simulation, (2024,), (*BENEFITS, *DOWNSTREAM)) extracted = _snapshot(reloaded, (2024,), (*BENEFITS, *DOWNSTREAM)) - assert extracted == {key: values[:1] for key, values in live.items()} + expected = {key: values[:1] for key, values in live.items()} + assert extracted == expected, { + key: {"live": expected[key], "reloaded": extracted[key]} + for key in expected + if extracted[key] != expected[key] + } result[f"extraction_reload_{label}"] = extracted + live_future = _snapshot(simulation, (2026, 2027), (*BENEFITS, *DOWNSTREAM)) + result[f"same_period_override_{label}"] = live_future + assert ( + live_future["medicare_irmaa_magi_two_years_prior@2026"] + == live["medicare_irmaa_tax_year_magi@2024"] + ) + extracted_future = _snapshot(reloaded, (2026, 2027), (*BENEFITS, *DOWNSTREAM)) + expected_future = {key: values[:1] for key, values in live_future.items()} + assert extracted_future == expected_future, { + key: {"live": expected_future[key], "reloaded": extracted_future[key]} + for key in expected_future + if extracted_future[key] != expected_future[key] + } + result[f"extraction_future_reload_{label}"] = extracted_future + reloaded.set_input("social_security_retirement", 2024, [36_000]) + result[f"extraction_rewrite_{label}"] = _snapshot( + reloaded, (2024, 2026, 2027), (*BENEFITS, *DOWNSTREAM) + ) del simulation, reloaded gc.collect() @@ -402,11 +436,18 @@ def differential(tmp_path_factory): "later_override_microsimulation", "extraction_reload_simulation", "extraction_reload_microsimulation", + "same_period_override_simulation", + "same_period_override_microsimulation", + "extraction_future_reload_simulation", + "extraction_future_reload_microsimulation", + "extraction_rewrite_simulation", + "extraction_rewrite_microsimulation", "automatic_baseline_simulation", "automatic_baseline_microsimulation", "neutral_constructor_simulation", "neutral_constructor_microsimulation", "complete_export_reload", + "complete_export_rewrite", "yaml_core_builder", "yaml_core_downstream", "public_set_input", diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index 49fe6da84b4..5ed5ef625b4 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -134,9 +134,10 @@ def test_later_public_override_survives_clones_and_cache_clear(wrapper): copy.set_input(public, 2024, [42]) drop_inherited_values(copy) assert amount(copy, public, 2024) == 42 + assert amount(copy, f"{public}_reported", 2024) == 42 assert simulation._user_input_keys == original_keys assert amount(simulation, public, 2024) == 0 - assert amount(simulation, f"{public}_reported", 2024) == 12_000 + assert amount(simulation, f"{public}_reported", 2024) == 0 @pytest.mark.parametrize( @@ -174,12 +175,17 @@ def test_core_builder_migrates_initial_public_inputs_and_preserves_precedence(ro assert "age" in person simulation.set_input(retirement, 2024, [0]) assert amount(simulation, retirement, 2024) == 0 - assert amount(simulation, f"{retirement}_reported", 2024) == 12_000 + assert amount(simulation, f"{retirement}_reported", 2024) == 0 + simulation.set_input(disability, 2024, [42]) + assert amount(simulation, disability, 2024) == 42 + assert amount(simulation, f"{disability}_reported", 2024) == 666 @pytest.mark.parametrize("wrapper", (*WRAPPERS, make_core_simulation)) @pytest.mark.parametrize("reported_input", (None, 0, 6_000)) -def test_later_public_write_seeds_only_absent_reported_input(wrapper, reported_input): +def test_later_public_writes_refresh_automatic_sources_and_keep_explicit_twins( + wrapper, reported_input +): public = SOCIAL_SECURITY_COMPONENTS[0] reported = f"{public}_reported" inputs = {} if reported_input is None else {reported: {2024: reported_input}} @@ -201,6 +207,66 @@ def test_later_public_write_seeds_only_absent_reported_input(wrapper, reported_i assert amount(simulation, public, 2026) == pytest.approx( expected_source * 13_522.16 / 12_000, abs=0.01 ) + simulation.set_input(public, 2024, [24_000]) + expected_source = 24_000 if reported_input is None else reported_input + assert amount(simulation, reported, 2024) == expected_source + drop_inherited_values(simulation) + assert amount(simulation, public, 2026) == pytest.approx( + expected_source * 13_522.16 / 12_000, abs=0.01 + ) + assert input_periods(simulation, reported) == {"2024"} + + +@pytest.mark.parametrize("wrapper", (*WRAPPERS, make_core_simulation)) +def test_repeated_public_writes_refresh_migrated_sources_and_future_years(wrapper): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + inputs = {public: {2024: 12_000}} + simulation = ( + wrapper(inputs) + if wrapper is make_core_simulation + else make_simulation(wrapper, inputs) + ) + simulation.set_input(FACTOR, 2024, [1.08]) + simulation.set_input(FACTOR, 2026, [1.08]) + for override in (24_000, 18_000, 0): + simulation.set_input(public, 2024, [override]) + drop_inherited_values(simulation) + assert amount(simulation, public, 2024) == override + assert amount(simulation, reported, 2024) == override + assert amount(simulation, public, 2026) == pytest.approx( + override * 13_522.16 / 12_000 * 1.08, abs=0.01 + ) + assert input_periods(simulation, public) == {"2024"} + assert input_periods(simulation, reported) == {"2024"} + + +@pytest.mark.parametrize("wrapper", (*WRAPPERS, make_core_simulation)) +@pytest.mark.parametrize("explicit_source", (0, 12_000)) +def test_direct_reported_holder_write_replaces_automatic_source_origin( + wrapper, explicit_source +): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + inputs = {public: {2024: 12_000}} + simulation = ( + wrapper(inputs) + if wrapper is make_core_simulation + else make_simulation(wrapper, inputs) + ) + # A same-value write also becomes explicit; provenance cannot be inferred + # by comparing the stored amount with its former automatic amount. + simulation.get_holder(reported).set_input( + 2024, [explicit_source], simulation.branch_name + ) + for override in (24_000, 18_000): + simulation.set_input(public, 2024, [override]) + drop_inherited_values(simulation) + assert amount(simulation, public, 2024) == override + assert amount(simulation, reported, 2024) == explicit_source + assert amount(simulation, public, 2026) == pytest.approx( + explicit_source * 13_522.16 / 12_000, abs=0.01 + ) @pytest.mark.parametrize("wrapper", (*WRAPPERS, make_core_simulation)) @@ -216,7 +282,9 @@ def test_later_monthly_public_write_seeds_the_stored_rolling_year(wrapper): assert input_periods(simulation, reported) == {"year:2024-07"} assert simulation.get_array(public, "year:2024-07")[0] == 12_000 assert simulation.get_array(reported, "year:2024-07")[0] == 12_000 - assert amount(simulation, public, "2024-07") == 1_000 + # A monthly output reads the calendar year, whose January start precedes + # the first rolling source in July. Main also returns zero here. + assert amount(simulation, public, "2024-07") == 0 assert amount(simulation, public, 2026) > 12_000 @@ -245,6 +313,54 @@ def test_clone_and_branch_later_public_writes_keep_input_lists_private(wrapper): assert input_periods(simulation, reported) == {"2024"} +@pytest.mark.parametrize("wrapper", WRAPPERS) +def test_branch_explicit_twin_overrides_inherited_automatic_source_privately(wrapper): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + simulation = make_simulation(wrapper, {public: {2024: 12_000}}) + explicit = simulation.get_branch("explicit_reported_source") + explicit.set_input(reported, 2024, [0]) + automatic = simulation.get_branch("inherited_automatic_source") + for copy in (explicit, automatic): + copy.set_input(public, 2024, [24_000]) + drop_inherited_values(copy) + assert amount(copy, public, 2024) == 24_000 + assert amount(explicit, reported, 2024) == 0 + assert amount(explicit, public, 2026) == 0 + assert amount(automatic, reported, 2024) == 24_000 + assert amount(automatic, public, 2026) == pytest.approx(27_044.32, abs=0.01) + assert amount(simulation, reported, 2024) == 12_000 + # A branch's explicit replacement must not mark its parent's automatic + # source explicit, or prevent the parent's next public write from seeding. + simulation.set_input(public, 2024, [36_000]) + assert amount(simulation, reported, 2024) == 36_000 + assert amount(explicit, reported, 2024) == 0 + assert amount(automatic, reported, 2024) == 24_000 + + +@pytest.mark.parametrize("wrapper", WRAPPERS) +def test_neutral_reform_replay_keeps_automatic_source_origin_private(wrapper): + from policyengine_core.reforms import Reform + + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + simulation = make_simulation(wrapper, {public: {2024: 12_000}}) + parent_later = amount(simulation, public, 2026) + branch = simulation.get_branch("replayed_automatic_source", clone_system=True) + branch.apply_reform( + Reform.from_dict( + {"gov.contrib.ssa.benefit_levers.across_the_board.rate": {"2024": 0}} + ) + ) + branch.set_input(public, 2024, [24_000]) + assert amount(branch, reported, 2024) == 24_000 + assert amount(branch, public, 2024) == 24_000 + assert amount(branch, public, 2026) == pytest.approx(27_044.32, abs=0.01) + assert input_periods(branch, reported) == {"2024"} + assert amount(simulation, reported, 2024) == 12_000 + assert amount(simulation, public, 2026) == parent_later + + def test_dataset_constructor_migrates_stored_original_columns(): public = SOCIAL_SECURITY_COMPONENTS[0] frame = pd.DataFrame( @@ -284,12 +400,12 @@ def test_dataset_constructor_migrates_stored_original_columns(): # Loader-time public writes migrate; later writes remain final outputs. simulation.set_input(public, 2024, [0]) assert amount(simulation, public, 2024) == 0 - assert amount(simulation, f"{public}_reported", 2024) == 12_000 + assert amount(simulation, f"{public}_reported", 2024) == 0 reloaded = reload_source_inputs( export_source_inputs(simulation), tax_benefit_system=system ) assert amount(reloaded, public, 2024) == 0 - assert amount(reloaded, f"{public}_reported", 2024) == 12_000 + assert amount(reloaded, f"{public}_reported", 2024) == 0 def test_complete_snapshot_preserves_histories_and_separate_final_overrides(): @@ -320,6 +436,10 @@ def test_complete_snapshot_preserves_histories_and_separate_final_overrides(): ) assert input_periods(reloaded, public) == {"2026", "2027"} assert export_source_inputs(reloaded)["overrides"][public]["2027"] == [0] + reloaded.set_input(public, 2027, [24_000]) + assert amount(reloaded, f"{public}_reported", 2027) == 24_000 + assert amount(reloaded, public, 2028) > 24_000 + assert amount(simulation, f"{public}_reported", 2027) == 0 def test_legacy_extraction_uses_reported_sources_and_months_keep_flow_conversion(): @@ -336,6 +456,9 @@ def test_legacy_extraction_uses_reported_sources_and_months_keep_flow_conversion assert amount(copy, public, "2024-07") == 1_000 assert amount(copy, public, 2024) == 12_000 assert amount(copy, public, "2024-07") == 1_000 + reloaded.set_input(public, 2024, [24_000]) + assert amount(reloaded, f"{public}_reported", 2024) == 24_000 + assert amount(reloaded, public, 2026) == pytest.approx(27_044.32, abs=0.01) @pytest.mark.parametrize("wrapper", WRAPPERS) @@ -348,10 +471,94 @@ def test_legacy_extraction_preserves_later_public_override(wrapper, override): extracted = simulation.extract_person(0) person = next(iter(extracted["people"].values())) assert person[public]["2024"] == override - assert person[reported]["2024"] == 12_000 + assert person[reported]["2024"] == override reloaded = Simulation(situation=extracted, tax_benefit_system=system) assert amount(reloaded, public, 2024) == override - assert amount(reloaded, reported, 2024) == 12_000 + assert amount(reloaded, reported, 2024) == override + + +@pytest.mark.parametrize("wrapper", WRAPPERS) +@pytest.mark.parametrize("snapshot", ("legacy_extraction", "complete_export")) +def test_source_round_trip_retains_automatic_origin_for_subsequent_writes( + wrapper, snapshot +): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + simulation = make_simulation(wrapper, {public: {2024: 12_000}}) + simulation.set_input(public, 2024, [24_000]) + # A cached intervening source is not exported or treated as supplied. + amount(simulation, reported, 2025) + if snapshot == "legacy_extraction": + payload = json.loads(json.dumps(simulation.extract_person(0))) + person = next(iter(payload["people"].values())) + reloaded = Simulation(situation=payload, tax_benefit_system=system) + assert person[public] == {"2024": 24_000} + else: + payload = json.loads(json.dumps(export_source_inputs(simulation))) + person = next(iter(payload["sources"]["people"].values())) + reloaded = reload_source_inputs(payload, tax_benefit_system=system) + assert public not in person + assert payload["overrides"][public] == {"2024": [24_000]} + assert person[reported] == {"2024": 24_000} + assert set(person).issubset(system.variables) + reloaded.set_input(public, 2024, [36_000]) + assert amount(reloaded, public, 2024) == 36_000 + assert amount(reloaded, reported, 2024) == 36_000 + assert amount(reloaded, public, 2026) == pytest.approx(40_566.48, abs=0.02) + assert input_periods(reloaded, reported) == {"2024"} + assert amount(simulation, reported, 2024) == 24_000 + + +@pytest.mark.parametrize("wrapper", WRAPPERS) +def test_complete_reload_reform_baseline_retains_automatic_source_origin(wrapper): + from policyengine_core.reforms import Reform + + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + simulation = make_simulation(wrapper, {public: {2024: 12_000}}) + simulation.set_input(public, 2024, [24_000]) + reloaded = reload_source_inputs( + export_source_inputs(simulation), + tax_benefit_system=system, + reform=Reform.from_dict( + {"gov.contrib.ssa.benefit_levers.across_the_board.rate": {"2024": 0}} + ), + ) + baseline = reloaded.baseline + assert baseline is not None + baseline.set_input(public, 2024, [36_000]) + assert amount(baseline, public, 2024) == 36_000 + assert amount(baseline, reported, 2024) == 36_000 + assert amount(baseline, public, 2026) == pytest.approx(40_566.48, abs=0.02) + assert amount(reloaded, reported, 2024) == 24_000 + reloaded.set_input(public, 2024, [48_000]) + assert amount(reloaded, reported, 2024) == 48_000 + assert amount(reloaded, public, 2026) == pytest.approx(54_088.64, abs=0.02) + assert amount(baseline, reported, 2024) == 36_000 + assert amount(simulation, reported, 2024) == 24_000 + + +@pytest.mark.parametrize("wrapper", WRAPPERS) +def test_explicit_only_legacy_extraction_needs_no_provenance_envelope(wrapper): + retirement, disability = SOCIAL_SECURITY_COMPONENTS[:2] + simulation = make_simulation( + wrapper, + { + f"{retirement}_reported": {2024: 0}, + f"{disability}_reported": {2024: 6_000}, + }, + ) + payload = simulation.extract_person(0) + assert set(payload).issubset(system.entities_plural()) + person = next(iter(payload["people"].values())) + assert set(person).issubset(system.variables) + reloaded = Simulation(situation=payload, tax_benefit_system=system) + reloaded.set_input(retirement, 2024, [24_000]) + reloaded.set_input(disability, 2024, [12_000]) + assert amount(reloaded, f"{retirement}_reported", 2024) == 0 + assert amount(reloaded, f"{disability}_reported", 2024) == 6_000 + assert amount(reloaded, retirement, 2026) == 0 + assert amount(reloaded, disability, 2026) == pytest.approx(6_761.08, abs=0.01) def test_reported_twins_are_pure_and_original_dataset_names_keep_their_uprater(): @@ -467,7 +674,7 @@ def test_exact_multi_year_sources_keep_keys_uniform_factors_and_overrides(wrappe export_source_inputs(clone), tax_benefit_system=system ) assert amount(override_reload, public, multi_year) == 0 - assert amount(override_reload, reported, multi_year) == 24_000 + assert amount(override_reload, reported, multi_year) == 0 assert input_periods(override_reload, "household_weight") == {multi_year} diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 3f94ce1c977..5e5a880fef5 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -1,13 +1,15 @@ """Compatibility ingress and complete, source-safe simulation snapshots. Initial ingress moves legacy benefit inputs to their pure reported sources. -A later public write overrides its period and seeds an absent reported source. +A later public write overrides its period and refreshes its automatic source. +An explicitly supplied reported source retains precedence. """ import json from contextvars import ContextVar from functools import wraps +from policyengine_core.periods import period as to_period from policyengine_core.simulations.simulation import NpEncoder from policyengine_us.tools.period_branch import delete_input_values @@ -19,11 +21,90 @@ "social_security_survivors", "social_security_dependents", ) +_REPORTED_COMPONENTS = tuple(f"{name}_reported" for name in SOCIAL_SECURITY_COMPONENTS) _initial_social_security_ingress = ContextVar( "initial_social_security_ingress", default=False ) +_PROVENANCE = "_social_security_input_provenance" +_AUTOMATIC_INPUTS = "_social_security_automatic_reported_inputs" + + +def _recorded_input_location(simulation, name, period, *, branch_names=None): + """Identify the actual supplied tier, including one hidden by a cache.""" + holder = simulation.get_holder(name) + period = holder._storage_period(period) + input_keys = getattr(simulation, "_user_input_keys", ()) + if branch_names is None: + branch_names = simulation._get_visible_branch_names() + for branch_name in branch_names: + key = (name, branch_name, period) + if key not in input_keys: + continue + for storage_name, storage, separator in ( + ("memory", holder._memory_storage, ":"), + ("disk", holder._disk_storage, "_"), + ): + if holder._stores_user_input(period, branch_name, storage_name): + sequence = storage._sequence_numbers.get( + f"{branch_name}{separator}{period}" + ) + return key, (storage_name, sequence) + return None + + +def _is_automatic_reported_input(simulation, name, period): + location = _recorded_input_location(simulation, name, period) + return ( + location is not None + and getattr(simulation, _AUTOMATIC_INPUTS, {}).get(location[0]) == location[1] + ) + + +def _mark_automatic_reported_input(simulation, name, period, *, branch_name=None): + location = _recorded_input_location( + simulation, + name, + period, + branch_names=None if branch_name is None else (branch_name,), + ) + if location is not None: + # Core clones simulation attributes shallowly. Never mutate a shared + # provenance dictionary: each write belongs to this simulation alone. + automatic = dict(getattr(simulation, _AUTOMATIC_INPUTS, {})) + automatic[location[0]] = location[1] + setattr(simulation, _AUTOMATIC_INPUTS, automatic) + + +def _automatic_reported_metadata(simulation, *, exported_periods=None): + automatic = {} + for public_name in SOCIAL_SECURITY_COMPONENTS: + name = f"{public_name}_reported" + for period in simulation._get_set_input_periods(name): + if ( + exported_periods is not None + and (name, str(period)) not in exported_periods + ): + continue + if _is_automatic_reported_input(simulation, name, period): + automatic.setdefault(name, []).append(str(period)) + return {"version": 1, "automatic_reported": automatic} if automatic else None + + +def _restore_automatic_reported_metadata(simulation, metadata): + if metadata is None: + return + if not isinstance(metadata, dict) or metadata.get("version") != 1: + raise ValueError("Unsupported Social Security input provenance version") + automatic = metadata.get("automatic_reported", {}) + if not isinstance(automatic, dict): + raise ValueError("Social Security input provenance must contain a mapping") + for name, periods in automatic.items(): + if name not in _REPORTED_COMPONENTS or not isinstance(periods, list): + raise ValueError("Invalid automatic Social Security reported input") + for period in periods: + _mark_automatic_reported_input(simulation, name, to_period(period)) def install_social_security_input_compatibility(): @@ -55,10 +136,10 @@ def set_input(simulation, variable_name, period, value): if not track_public_write: return original_set_input(simulation, variable_name, period, value) - holder = simulation.get_holder(variable_name) - branch = simulation.branch_name before = { - supplied_period: holder._stored_sequence_number(supplied_period, branch) + supplied_period: _recorded_input_location( + simulation, variable_name, supplied_period + ) for supplied_period in simulation._get_set_input_periods(variable_name) } result = original_set_input(simulation, variable_name, period, value) @@ -68,10 +149,13 @@ def set_input(simulation, variable_name, period, value): # in: a monthly input to a yearly variable stores a rolling year. # Inspect actual stores rather than guessing those period aliases. for supplied_period in simulation._get_set_input_periods(variable_name): - if ( - holder._stored_sequence_number(supplied_period, branch) - == before.get(supplied_period) - or supplied_period in reported_periods + if _recorded_input_location( + simulation, variable_name, supplied_period + ) == before.get(supplied_period) or ( + supplied_period in reported_periods + and not _is_automatic_reported_input( + simulation, reported_name, supplied_period + ) ): continue supplied_value = simulation._get_recorded_input_array( @@ -84,6 +168,9 @@ def set_input(simulation, variable_name, period, value): supplied_period, supplied_value.copy(), ) + _mark_automatic_reported_input( + simulation, reported_name, supplied_period + ) # Core clones share this legacy list. Replace it privately so a # later override is extracted without changing the source clone. @@ -97,6 +184,65 @@ def set_input(simulation, variable_name, period, value): set_input._social_security_compatibility = True CoreSimulation.set_input = set_input + original_extract = CoreSimulation.extract_person + if not getattr(original_extract, "_social_security_compatibility", False): + + @wraps(original_extract) + def extract_person(simulation, index=0, exclude_entities=("state",)): + payload = original_extract(simulation, index, exclude_entities) + if not getattr( + simulation.tax_benefit_system, + "_social_security_input_compatibility", + False, + ): + return payload + people = payload.get(simulation.persons.entity.plural, {}) + exported_periods = { + (name, str(period)) + for person in people.values() + for name, values in person.items() + if name in _REPORTED_COMPONENTS + for period in values + } + metadata = _automatic_reported_metadata( + simulation, exported_periods=exported_periods + ) + if metadata is not None: + payload[_PROVENANCE] = metadata + return payload + + extract_person._social_security_compatibility = True + CoreSimulation.extract_person = extract_person + + original_invalidate = CoreSimulation._invalidate_all_caches + if not getattr(original_invalidate, "_social_security_compatibility", False): + + @wraps(original_invalidate) + def invalidate_all_caches(simulation): + automatic = [] + if getattr( + simulation.tax_benefit_system, + "_social_security_input_compatibility", + False, + ): + for key, stamp in getattr(simulation, _AUTOMATIC_INPUTS, {}).items(): + name, branch_name, period = key + if _recorded_input_location( + simulation, name, period, branch_names=(branch_name,) + ) == (key, stamp): + automatic.append(key) + result = original_invalidate(simulation) + # Core replay assigns new store sequence numbers. Replaying an + # automatic source does not turn it into an explicit caller input. + for name, branch_name, period in automatic: + _mark_automatic_reported_input( + simulation, name, period, branch_name=branch_name + ) + return result + + invalidate_all_caches._social_security_compatibility = True + CoreSimulation._invalidate_all_caches = invalidate_all_caches + def wrap_builder(original_build): @wraps(original_build) def build(builder, tax_benefit_system, input_dict, simulation=None): @@ -106,6 +252,13 @@ def build(builder, tax_benefit_system, input_dict, simulation=None): return original_build( builder, tax_benefit_system, input_dict, simulation ) + metadata = input_dict.get(_PROVENANCE) + if _PROVENANCE in input_dict: + input_dict = { + name: values + for name, values in input_dict.items() + if name != _PROVENANCE + } token = _initial_social_security_ingress.set(True) try: result = original_build( @@ -135,12 +288,17 @@ def build(builder, tax_benefit_system, input_dict, simulation=None): ) ) migrate_social_security_inputs(result) + _restore_automatic_reported_metadata(result, metadata) return result build._social_security_compatibility = True return build - for method_name in ("build_from_entities", "build_from_variables"): + for method_name in ( + "build_from_dict", + "build_from_entities", + "build_from_variables", + ): original_build = getattr(SimulationBuilder, method_name) if not getattr(original_build, "_social_security_compatibility", False): setattr(SimulationBuilder, method_name, wrap_builder(original_build)) @@ -190,6 +348,7 @@ def migrate_social_security_inputs(simulation, *, include_baseline=True): break for period, value in sources.items(): simulation.set_input(reported_name, period, value.copy()) + _mark_automatic_reported_input(simulation, reported_name, period) for period in sources: for branch_name in visible_branches: delete_input_values(holder, period, branch_name) @@ -276,6 +435,7 @@ def export_source_inputs(simulation): "version": 1, "sources": sources, "overrides": overrides, + _PROVENANCE: _automatic_reported_metadata(simulation), "default_input_period": str(simulation.default_input_period), "default_calculation_period": str( simulation.default_calculation_period @@ -299,7 +459,14 @@ def reload_source_inputs(payload, **simulation_kwargs): raise ValueError("Unsupported source-input snapshot version") for period_name in ("default_input_period", "default_calculation_period"): simulation_kwargs.setdefault(period_name, payload[period_name]) - simulation = Simulation(situation=payload["sources"], **simulation_kwargs) + metadata = payload.get(_PROVENANCE) + sources = payload["sources"] + if metadata is not None: + # Restore before Core creates a reform baseline, so that clone also + # knows which twins came from legacy public inputs. + sources = {**sources, _PROVENANCE: metadata} + simulation = Simulation(situation=sources, **simulation_kwargs) + _restore_automatic_reported_metadata(simulation, metadata) for name, periods in payload["overrides"].items(): for period, values in periods.items(): simulation.set_input(name, period, values) From d0187abd9541156a3a7e6b125eeb53efdc2e9167 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 17:21:11 -0400 Subject: [PATCH 20/76] Require an independent main reference and nonzero dataset evidence --- .../test_social_security_current_law.py | 62 +++++++++++++++++-- 1 file changed, 56 insertions(+), 6 deletions(-) diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 0dbd6cc0190..3798b71a342 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -2,8 +2,9 @@ The same probe runs in separate Python processes against this checkout and main. No variable definitions or input names are substituted in the reference. -Set PE_US_MAIN_CHECKOUT to a clean main checkout, or make a local main remote -ref available (the fixture archives it without changing either checkout). +An unmodified pre-feature main checkout is required. Set PE_US_MAIN_CHECKOUT +to that checkout, or make the reviewed main commit or a pre-feature main ref +available locally (the fixture archives it without changing either checkout). """ import io @@ -45,6 +46,7 @@ "medicare_part_b_premium", ) AMOUNTS = ((12_000, 0, 1_800, 600), (0, 6_000, 1_200, 300)) +REVIEWED_MAIN_SHA = "8986e05cab65d56435a23df1910e2d35f2c5f5bf" def situation(*, include_benefits=True): @@ -156,6 +158,8 @@ class HouseholdMicrosimulation(Microsimulation): reload_source_inputs, ) except ImportError: + if has_levers: + raise result["complete_export_reload"] = _snapshot( simulation, (2024, 2025, 2026, 2027) ) @@ -321,6 +325,12 @@ class HouseholdMicrosimulation(Microsimulation): result["h5_scored_year"] = _snapshot( simulation, (2026,), (*BENEFITS, *DOWNSTREAM) ) + for index, component in enumerate(COMPONENTS): + np.testing.assert_array_equal( + np.asarray(simulation.calculate(component, 2026)), + [amounts[index] for amounts in AMOUNTS], + err_msg=f"H5 scored-year source was not loaded for {component}", + ) # Main has no benefit-lever parameter. Its role here is the unchanged # current-law reference; the lever assertion runs in the current tree. if has_levers: @@ -352,9 +362,13 @@ class HouseholdMicrosimulation(Microsimulation): def _main_checkout(current, directory): + reported_files = {f"{component}_reported.py" for component in COMPONENTS} configured = os.environ.get("PE_US_MAIN_CHECKOUT") if configured: checkout = Path(configured).resolve() + assert checkout != current.resolve(), ( + "PE_US_MAIN_CHECKOUT must be a separate unmodified pre-feature main checkout" + ) assert (checkout / "policyengine_us").is_dir(), configured if (checkout / ".git").exists(): status = subprocess.check_output( @@ -363,16 +377,51 @@ def _main_checkout(current, directory): text=True, ) assert not status, "PE_US_MAIN_CHECKOUT must have no tracked changes" + if (current / ".git").exists(): + current_head = subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=current, text=True + ).strip() + reference_head = subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=checkout, text=True + ).strip() + assert reference_head != current_head, ( + "PE_US_MAIN_CHECKOUT must not use the current checkout's HEAD" + ) + assert not any( + path.name in reported_files + for path in (checkout / "policyengine_us" / "variables").rglob( + "*_reported.py" + ) + ), "PE_US_MAIN_CHECKOUT must contain unmodified pre-feature main variables" return checkout - for ref in ("refs/remotes/origin/main", "refs/remotes/gh/main", "refs/heads/main"): + for ref in ( + REVIEWED_MAIN_SHA, + "refs/remotes/origin/main", + "refs/remotes/gh/main", + "refs/heads/main", + ): exists = subprocess.run( - ["git", "rev-parse", "--verify", ref], + ["git", "rev-parse", "--verify", f"{ref}^{{commit}}"], cwd=current, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=False, ) if exists.returncode == 0: + names = subprocess.check_output( + [ + "git", + "ls-tree", + "-r", + "--name-only", + ref, + "policyengine_us/variables", + ], + cwd=current, + text=True, + ) + if any(Path(name).name in reported_files for name in names.splitlines()): + continue archive = subprocess.check_output(["git", "archive", ref], cwd=current) checkout = directory / "main" checkout.mkdir() @@ -380,8 +429,9 @@ def _main_checkout(current, directory): stored.extractall(checkout, filter="data") return checkout pytest.skip( - "Differential needs PE_US_MAIN_CHECKOUT or a local origin/main, gh/main, " - "or main ref containing an unmodified main checkout" + "Differential needs an unmodified pre-feature main checkout: set " + "PE_US_MAIN_CHECKOUT or provide the reviewed main commit or a local " + "origin/main, gh/main, or main ref without reported benefit variables" ) From 9bb01d88141d93b670dbe6d4afdd5cb3637f4776 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 17:24:21 -0400 Subject: [PATCH 21/76] Require compatible Core input provenance APIs --- docs/social-security-benefit-levers.md | 3 +++ pyproject.toml | 2 +- uv.lock | 8 ++++---- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/social-security-benefit-levers.md b/docs/social-security-benefit-levers.md index 9b3fd078421..ca76fa723d4 100644 --- a/docs/social-security-benefit-levers.md +++ b/docs/social-security-benefit-levers.md @@ -42,6 +42,9 @@ explicit sources through legacy extraction and complete snapshots, including subsequent public writes after reload. Builders consume this metadata before validating ordinary situation inputs. +This compatibility layer requires Core 3.33.1 or later for recorded-input, +storage-origin and replay APIs. The lockfile uses Core 3.33.2. + Each component has a `_source` selector with default `reported`. Reserved options `supplied_aime`, `supplied_pia`, `supplied_ratio_survivor`, and `history` fail explicitly until their calculation routes are implemented. Selection never diff --git a/pyproject.toml b/pyproject.toml index fe70927a79d..4087f294fde 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ dependencies = [ # pandas 3 copy-on-write enables the shallow-copy fast path in dataset # extension and is available on every supported Python version. "pandas>=3.0", - "policyengine-core>=3.32.27", + "policyengine-core>=3.33.1", "spm-calculator>=1.0.0,<=1.0.0.post1", "tables>=3.9", "tqdm>=4.67.1", diff --git a/uv.lock b/uv.lock index 42a51471367..9ce7bfdec91 100644 --- a/uv.lock +++ b/uv.lock @@ -1786,7 +1786,7 @@ wheels = [ [[package]] name = "policyengine-core" -version = "3.32.27" +version = "3.33.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "dpath" }, @@ -1806,9 +1806,9 @@ dependencies = [ { name = "standard-imghdr" }, { name = "wheel" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/de/24/84c0747c64fba95e89d3acd332604ad01e8a220ce7e80cf8d342b3b213bd/policyengine_core-3.32.27.tar.gz", hash = "sha256:632e8708ec243f7b592ec857e1978df207547c1510039373e707c310173cf4e4", size = 599334, upload-time = "2026-10-09T05:44:09.236Z" } +sdist = { url = "https://files.pythonhosted.org/packages/84/dc/b6e44d9cdb9afee354f89584eebe310d2e6480f852ca51d55570eee5a648/policyengine_core-3.33.2.tar.gz", hash = "sha256:4f0420021a5439b8d3b7c86333e291b003f4be689aebfff49cb10ad1b6952ccb", size = 693096, upload-time = "2026-10-10T19:09:47.182Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/cc/b3/fc1a10b03f2a159bc5623f0aa465887d1c1c98ee2f3118fd9ee3e8275921/policyengine_core-3.32.27-py3-none-any.whl", hash = "sha256:d0b0135939252ed628efd496988f76c0ab89bfa3b3437aa90f39000667c37687", size = 278496, upload-time = "2026-10-09T05:44:07.764Z" }, + { url = "https://files.pythonhosted.org/packages/43/1b/9950def57dc9ec82aa2647aed02237199e347b5b7ecb898d5dfaf70014bc/policyengine_core-3.33.2-py3-none-any.whl", hash = "sha256:631e20811b74f6a7df906891efa5c78fbfd461ecefd65902b6e1456b1a0e4098", size = 301904, upload-time = "2026-10-10T19:09:45.432Z" }, ] [[package]] @@ -1845,7 +1845,7 @@ requires-dist = [ { name = "jupyter-book", marker = "extra == 'dev'", specifier = ">=1.0.4.post1" }, { name = "microdf-python", specifier = ">=1.0.0" }, { name = "pandas", specifier = ">=3.0" }, - { name = "policyengine-core", specifier = ">=3.32.27" }, + { name = "policyengine-core", specifier = ">=3.33.1" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.9.0" }, { name = "setuptools", marker = "extra == 'dev'", specifier = ">=80.9.0" }, { name = "spm-calculator", specifier = ">=1.0.0,<=1.0.0.post1" }, From 4372e6bcff639cfbf5bebaf7904106cf08094a9b Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 17:54:00 -0400 Subject: [PATCH 22/76] Isolate extraction comparisons with the review reproducer --- .../test_social_security_current_law.py | 38 +++++++++++++++---- 1 file changed, 30 insertions(+), 8 deletions(-) diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 3798b71a342..1aed74b2481 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -74,6 +74,24 @@ def situation(*, include_benefits=True): return result +def extraction_situation(): + """The review's single-person, single-input-period extraction reproducer.""" + result = { + "people": { + "retiree": { + "age": {2024: 70}, + "social_security_retirement": {2024: 12_000}, + "taxable_interest_income": {2024: 40_000}, + "takes_up_ssi_if_eligible": {2024: True}, + } + } + } + for plural in ("households", "tax_units", "spm_units", "families", "marital_units"): + result[plural] = {"group": {"members": ["retiree"]}} + result["households"]["group"]["state_code"] = {2024: "TX"} + return result + + def _snapshot(simulation, periods, variables=BENEFITS): return { f"{variable}@{period}": np.asarray( @@ -188,17 +206,15 @@ class HouseholdMicrosimulation(Microsimulation): # from public ingress. An explicitly supplied twin retains priority # in the separate precedence tests. Extraction must preserve this # origin as well as the final override for a subsequent public write. - inputs = situation() - for person in inputs["people"].values(): - # Legacy extraction does not round-trip the existing earnings - # constructor migrations. Isolate the reported benefit regression. - person.pop("employment_income") - person.pop("self_employment_income") - simulation = wrapper(situation=inputs, tax_benefit_system=system) - simulation.set_input("social_security_retirement", 2024, [24_000, 0]) + simulation = wrapper( + situation=extraction_situation(), tax_benefit_system=system + ) + simulation.set_input("social_security_retirement", 2024, [24_000]) # Extract before calculations introduce earlier lookback cache periods # into Core's legacy first-known-period selection. payload = simulation.extract_person(0) + extracted_person = next(iter(payload["people"].values())) + assert extracted_person["age"] == {"2024": 70} reloaded = Simulation(situation=payload, tax_benefit_system=system) live = _snapshot(simulation, (2024,), (*BENEFITS, *DOWNSTREAM)) extracted = _snapshot(reloaded, (2024,), (*BENEFITS, *DOWNSTREAM)) @@ -209,6 +225,12 @@ class HouseholdMicrosimulation(Microsimulation): if extracted[key] != expected[key] } result[f"extraction_reload_{label}"] = extracted + assert live["medicare_irmaa_tax_year_magi@2024"][0] > 0 + for candidate in (simulation, reloaded): + # PE-US carries the supplied age forward. Keep this retired + # household eligible for Medicare in the future lookback checks. + for year in (2026, 2027): + np.testing.assert_array_equal(candidate.calculate("age", year), [70]) live_future = _snapshot(simulation, (2026, 2027), (*BENEFITS, *DOWNSTREAM)) result[f"same_period_override_{label}"] = live_future assert ( From dfd2e17c73f5fb4edc6f92b2d2487c69197d5611 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 18:02:47 -0400 Subject: [PATCH 23/76] Compare every Core builder form with legacy main inputs --- .../test_social_security_current_law.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 1aed74b2481..37b41f366f3 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -281,6 +281,23 @@ class HouseholdMicrosimulation(Microsimulation): del simulation gc.collect() + simulation = SimulationBuilder().build_from_entities(system, situation()) + result["core_entities_builder"] = _snapshot(simulation, (2024, 2026)) + del simulation + gc.collect() + + for route, method in ( + ("core_variables_builder", "build_from_variables"), + ("core_flat_dict_builder", "build_from_dict"), + ): + inputs = {"age": {2024: [70, 45]}} + for index, component in enumerate(COMPONENTS): + inputs[component] = {2024: [amounts[index] for amounts in AMOUNTS]} + simulation = getattr(SimulationBuilder(), method)(system, inputs) + result[route] = _snapshot(simulation, (2024, 2026)) + del simulation + gc.collect() + simulation = SimulationBuilder().build_from_dict( system, situation(include_benefits=False) ) @@ -522,6 +539,9 @@ def differential(tmp_path_factory): "complete_export_rewrite", "yaml_core_builder", "yaml_core_downstream", + "core_entities_builder", + "core_variables_builder", + "core_flat_dict_builder", "public_set_input", "clone", "branch", From 1ed0b378bb1aaa2a93c2c835c7e67226c02937be Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 19:26:49 -0400 Subject: [PATCH 24/76] Fetch the independent main reference for CI comparisons --- .github/workflows/push.yaml | 2 + .github/workflows/python-tests.yaml | 2 + .../code_health/test_make_test_matches_ci.py | 39 +++++++++++++++++++ 3 files changed, 43 insertions(+) diff --git a/.github/workflows/push.yaml b/.github/workflows/push.yaml index 1456db83bec..75a0c06a26c 100644 --- a/.github/workflows/push.yaml +++ b/.github/workflows/push.yaml @@ -490,6 +490,8 @@ jobs: - name: Update system certificate authorities id: system-ca uses: ./.github/actions/setup-system-ca + - name: Fetch Social Security differential reference + run: git fetch --no-tags --depth=1 origin 8986e05cab65d56435a23df1910e2d35f2c5f5bf - name: Set up Python uses: actions/setup-python@v6 with: diff --git a/.github/workflows/python-tests.yaml b/.github/workflows/python-tests.yaml index 4e682b6ed72..7294c69a808 100644 --- a/.github/workflows/python-tests.yaml +++ b/.github/workflows/python-tests.yaml @@ -146,6 +146,8 @@ jobs: - name: Update system certificate authorities id: system-ca uses: ./.github/actions/setup-system-ca + - name: Fetch Social Security differential reference + run: git fetch --no-tags --depth=1 origin 8986e05cab65d56435a23df1910e2d35f2c5f5bf - name: Set up Python uses: actions/setup-python@v6 with: diff --git a/policyengine_us/tests/code_health/test_make_test_matches_ci.py b/policyengine_us/tests/code_health/test_make_test_matches_ci.py index c06788638a5..50658ba6e4b 100644 --- a/policyengine_us/tests/code_health/test_make_test_matches_ci.py +++ b/policyengine_us/tests/code_health/test_make_test_matches_ci.py @@ -133,3 +133,42 @@ def test_contrib_ssa_yaml_has_a_dedicated_single_file_serial_route(): excluded = arguments[arguments.index("--exclude") + 1].split(",") assert "ssa" in excluded, "SSA YAML must not also enter a shared subdir batch" assert commands <= make_commands("test"), "The full suite must retain SSA's route" + + +@pytest.mark.parametrize("workflow_name", ("python-tests.yaml", "push.yaml")) +def test_microsimulation_fetches_the_pinned_social_security_reference(workflow_name): + """A shallow CI checkout otherwise skips every legacy differential case.""" + source = ( + REPO_ROOT + / "policyengine_us/tests/microsimulation/test_social_security_current_law.py" + ).read_text() + reference = re.search( + r"^REVIEWED_MAIN_SHA = [\"']([0-9a-f]{40})[\"']$", source, re.MULTILINE + ) + assert reference, "the differential must declare its reviewed main commit" + expected_fetch = [ + "git", + "fetch", + "--no-tags", + "--depth=1", + "origin", + reference.group(1), + ] + workflow = yaml.safe_load( + (REPO_ROOT / ".github/workflows" / workflow_name).read_text() + ) + steps = workflow["jobs"]["Microsimulation"]["steps"] + fetch_steps = [ + index + for index, step in enumerate(steps) + for line in (step.get("run") or "").splitlines() + if shlex.split(line) == expected_fetch + ] + test_steps = [ + index + for index, step in enumerate(steps) + if "test-microsimulation" in MAKE_TARGET.findall(step.get("run") or "") + ] + assert len(fetch_steps) == 1, "fetch the exact pre-feature reference once" + assert len(test_steps) == 1, "the microsimulation suite must run once" + assert fetch_steps[0] < test_steps[0], "fetch the reference before the suite runs" From 90c64b57855a50abb1409c2808b6ed59fc80574f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 20:42:10 -0400 Subject: [PATCH 25/76] Preserve every recorded benefit override in household extraction --- docs/social-security-benefit-levers.md | 9 ++- .../test_social_security_inputs.py | 61 +++++++++++++++++++ .../tools/social_security_inputs.py | 35 ++++++++++- 3 files changed, 99 insertions(+), 6 deletions(-) diff --git a/docs/social-security-benefit-levers.md b/docs/social-security-benefit-levers.md index ca76fa723d4..935fcfe4611 100644 --- a/docs/social-security-benefit-levers.md +++ b/docs/social-security-benefit-levers.md @@ -57,9 +57,12 @@ factors once, and perform their earnings test on the resulting payable amount. exports all recorded pure-input periods and entity memberships as `sources`, with supplied formula-variable values in a separate `overrides` payload. Cached calculations are excluded. `reload_source_inputs(payload, ...)` restores sources -before final overrides. Core's legacy `extract_person` remains a one-period -interface. Neither interface settles the separate API/application metadata and -field-exposure contract. +before final overrides. Core's legacy `extract_person` retains every recorded +input period for the four benefit components and their reported twins, excluding +formula caches. This keeps later overrides final on reload and preserves their +source histories and future uprating. Other variables keep Core's usual first +known period. Neither interface settles the separate API/application metadata +and field-exposure contract. Annual benefit inputs retain Core's period normalization. For example, supplying `social_security_retirement: {"2024-07": 500}` through a PE-US constructor stores diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index 5ed5ef625b4..0a5f3c98135 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -477,6 +477,67 @@ def test_legacy_extraction_preserves_later_public_override(wrapper, override): assert amount(reloaded, reported, 2024) == override +@pytest.mark.parametrize("wrapper", WRAPPERS) +@pytest.mark.parametrize("explicit_source", (None, 0, 6_000)) +def test_legacy_extraction_keeps_later_override_and_matching_source( + wrapper, explicit_source +): + from policyengine_core.reforms import Reform + + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + inputs = situation({public: {2024: 1_000}}) + inputs["people"]["other"] = {"age": {2024: 70}, public: {2024: 12_000}} + for plural, groups in inputs.items(): + if plural != "people": + groups["other"] = {"members": ["other"]} + reform = Reform.from_dict( + { + "gov.contrib.ssa.benefit_levers.across_the_board.start_year": { + "2026": 2026 + }, + "gov.contrib.ssa.benefit_levers.across_the_board.rate": {"2026": 0.08}, + } + ) + simulation = wrapper(situation=inputs, tax_benefit_system=system, reform=reform) + # Core's first known public period is now a calculated 2024 cache, so it + # cannot identify the later supplied override without recorded-input APIs. + np.testing.assert_array_equal(simulation.calculate(public, 2024), [1_000, 12_000]) + if explicit_source is not None: + simulation.set_input(reported, 2026, [500, explicit_source]) + simulation.set_input(public, 2026, [18_000, 24_000]) + expected_source = 24_000 if explicit_source is None else explicit_source + + # Select the second household: its Core extraction label retains index 2, + # while its reloaded population contains just one person at index 0. + payload = simulation.extract_person(1) + assert set(payload["people"]) == {"person_2"} + person = payload["people"]["person_2"] + assert person[public] == {"2026": 24_000} + assert person[reported] == {"2024": 12_000, "2026": expected_source} + reloaded = Simulation(situation=payload, tax_benefit_system=system, reform=reform) + assert input_periods(reloaded, public) == {"2026"} + # Reload must retain the final override, despite the 1.08 benefit factor. + assert amount(reloaded, public, 2026) == 24_000 + assert amount(reloaded, reported, 2026) == expected_source + assert amount(reloaded, public, 2027) == pytest.approx( + np.asarray(simulation.calculate(public, 2027))[1] + ) + + reloaded.set_input(public, 2026, [36_000]) + assert amount(reloaded, reported, 2026) == ( + 36_000 if explicit_source is None else explicit_source + ) + assert np.asarray(simulation.calculate(reported, 2026))[1] == expected_source + simulation.set_input(public, 2026, [18_000, 36_000]) + for candidate in (simulation, reloaded): + drop_inherited_values(candidate) + assert amount(reloaded, public, 2026) == 36_000 + assert amount(reloaded, public, 2027) == pytest.approx( + np.asarray(simulation.calculate(public, 2027))[1] + ) + + @pytest.mark.parametrize("wrapper", WRAPPERS) @pytest.mark.parametrize("snapshot", ("legacy_extraction", "complete_export")) def test_source_round_trip_retains_automatic_origin_for_subsequent_writes( diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 5e5a880fef5..4593ccbdccd 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -197,6 +197,34 @@ def extract_person(simulation, index=0, exclude_entities=("state",)): ): return payload people = payload.get(simulation.persons.entity.plural, {}) + benefit_inputs = { + name: { + str(period): value + for period in simulation._get_set_input_periods(name) + if (value := simulation._get_recorded_input_array(name, period)) + is not None + } + for name in (*SOCIAL_SECURITY_COMPONENTS, *_REPORTED_COMPONENTS) + } + for person_name, person in people.items(): + # Core names extracted people by their original zero-based + # index plus one, even when extracting a later household. + person_index = ( + int(person_name.removeprefix(f"{simulation.persons.entity.key}_")) + - 1 + ) + for name, values in benefit_inputs.items(): + # Core exports its first known period, which may be a + # formula cache rather than the later public override. + # Keep every actual benefit input and its source history + # together, without converting calculated values to inputs. + if values: + person[name] = { + period: float(value[person_index]) + for period, value in values.items() + } + else: + person.pop(name, None) exported_periods = { (name, str(period)) for person in people.values() @@ -392,9 +420,10 @@ def export_source_inputs(simulation): payload; calculated results never become sources. Only the current branch and its visible ancestors are exported, using their normal precedence. - This is the complete-history alternative to Core's legacy one-period - ``extract_person`` API. Export the desired small simulation or household; - this helper deliberately does not materialize a population dataset. + This is the complete-history alternative to Core's legacy ``extract_person`` + API, which retains all supplied periods only for benefit components and their + reported twins. Export the desired small simulation or household; this + helper deliberately does not materialize a population dataset. """ names = { key: [str(name) for name in population.ids] From a8b1ec85ddc74a61596c386031e560788b2b9864 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 24 Oct 2026 08:56:33 -0400 Subject: [PATCH 26/76] Preserve legacy Social Security inputs across Core write routes --- policyengine_us/system.py | 7 +- .../test_social_security_current_law.py | 80 +++++- .../test_social_security_inputs.py | 98 ++++++-- .../tools/social_security_inputs.py | 55 +++- .../tools/social_security_periods.py | 235 ++++++++++-------- 5 files changed, 325 insertions(+), 150 deletions(-) diff --git a/policyengine_us/system.py b/policyengine_us/system.py index 6e82e2e0635..ebbcde84794 100644 --- a/policyengine_us/system.py +++ b/policyengine_us/system.py @@ -40,7 +40,6 @@ install_social_security_input_compatibility, migrate_social_security_inputs, ) -from .tools.social_security_periods import SocialSecurityPeriodMixin from .tools.section_911 import validate_section_911_batch_inputs from .tools.per_capita_uprating import ( add_per_capita_parameters_for_parameter_uprating, @@ -226,7 +225,7 @@ def _backfill_state_code_from_str(simulation): state_code_str.delete_arrays(known_period) -class Simulation(SocialSecurityPeriodMixin, SPMSimulationMixin, CoreSimulation): +class Simulation(SPMSimulationMixin, CoreSimulation): """ A simulation of the tax-benefit system for the United States, defined against the base simulation class in the -core package. @@ -429,9 +428,7 @@ def _is_hdfstore_format(file_path): return False -class Microsimulation( - SocialSecurityPeriodMixin, SPMSimulationMixin, CoreMicrosimulation -): +class Microsimulation(SPMSimulationMixin, CoreMicrosimulation): """ A microsimulation of the tax-benefit system for the United States, defined against the base microsimulation class in the -core package. diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 37b41f366f3..aad399d0e53 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -45,6 +45,12 @@ "medicare_irmaa_magi", "medicare_part_b_premium", ) +INGRESS_OUTCOMES = ( + *COMPONENTS, + "income_tax", + "household_net_income", + "medicare_irmaa_magi_two_years_prior", +) AMOUNTS = ((12_000, 0, 1_800, 600), (0, 6_000, 1_200, 300)) REVIEWED_MAIN_SHA = "8986e05cab65d56435a23df1910e2d35f2c5f5bf" @@ -124,16 +130,29 @@ def _probe(): """Only imports from the checkout selected through the subprocess path.""" import gc + import pandas as pd import yaml from policyengine_core.data import Dataset + from policyengine_core.periods import period from policyengine_core.reforms import Reform - from policyengine_core.simulations import SimulationBuilder + from policyengine_core.simulations import ( + Simulation as CoreSimulation, + SimulationBuilder, + ) from policyengine_us import Microsimulation, Simulation from policyengine_us.system import system class HouseholdMicrosimulation(Microsimulation): default_dataset = None + class HolderSetterSimulation(Simulation): + def set_input(self, variable, input_period, value): + if variable in COMPONENTS and hasattr(self, "input_variables"): + return self.get_holder(variable).set_input( + period(input_period), np.asarray(value), self.branch_name + ) + return super().set_input(variable, input_period, value) + has_levers = "social_security_retirement_reported" in system.variables neutral = ( Reform.from_dict( @@ -307,6 +326,61 @@ class HouseholdMicrosimulation(Microsimulation): del simulation gc.collect() + # Both paths bypass CoreSimulation.set_input. The public holder must + # still seed future benefits and keep a later write final in its year. + for route, wrapper in ( + ("direct_public_holder", Simulation), + ("subclass_public_set_input_holder", HolderSetterSimulation), + ): + inputs = extraction_situation() + del inputs["people"]["retiree"]["social_security_retirement"] + simulation = wrapper(situation=inputs, tax_benefit_system=system) + if route == "direct_public_holder": + simulation.get_holder("social_security_retirement").set_input( + period(2024), np.array([12_000]), simulation.branch_name + ) + else: + simulation.set_input("social_security_retirement", 2024, [12_000]) + result[route] = _snapshot(simulation, (2024, 2026), INGRESS_OUTCOMES) + del simulation + gc.collect() + + # Raw Core dataset construction writes to holders directly, without + # either US wrapper or the situation builder's migration path. + columns = {"person_id": [1]} + for entity in ("household", "tax_unit", "spm_unit", "family", "marital_unit"): + columns[f"person_{entity}_id"] = [1] + columns[f"{entity}_id"] = [1] + columns.update( + { + "age__2024": [70], + "social_security_retirement__2024": [12_000], + } + ) + simulation = CoreSimulation( + dataset=pd.DataFrame(columns), + tax_benefit_system=system, + default_input_period=2024, + default_calculation_period=2024, + ) + result["raw_core_dataset"] = _snapshot(simulation, (2024, 2026), INGRESS_OUTCOMES) + del simulation + gc.collect() + + # Preserve Core's exact compound input and its annual/downstream + # behavior. This simulation has no US compound-period adapter. + inputs = extraction_situation() + inputs["people"]["retiree"]["social_security_retirement"] = {"year:2024:2": 24_000} + simulation = SimulationBuilder().build_from_dict(system, inputs) + result["core_builder_exact_multiyear"] = _snapshot( + simulation, ("year:2024:2",), ("social_security_retirement",) + ) + result["core_builder_exact_multiyear"].update( + _snapshot(simulation, (2024, 2026), INGRESS_OUTCOMES) + ) + del simulation + gc.collect() + for route in ("clone", "branch"): simulation = Simulation(situation=situation(), tax_benefit_system=system) original_inputs = list(simulation.input_variables) @@ -543,6 +617,10 @@ def differential(tmp_path_factory): "core_variables_builder", "core_flat_dict_builder", "public_set_input", + "direct_public_holder", + "subclass_public_set_input_holder", + "raw_core_dataset", + "core_builder_exact_multiyear", "clone", "branch", "clone_parent", diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index 0a5f3c98135..d1b899193d3 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -10,7 +10,11 @@ import numpy as np import pandas as pd import pytest -from policyengine_core.simulations import SimulationBuilder +from policyengine_core.periods import period +from policyengine_core.simulations import ( + Simulation as CoreSimulation, + SimulationBuilder, +) from policyengine_us import Microsimulation, Simulation from policyengine_us.data.dataset_schema import USSingleYearDataset @@ -71,6 +75,29 @@ def make_core_simulation(inputs=None, route="dict_entities"): return builder.build_from_variables(system, variables) +WRAPPER_WRITE_ROUTES = ( + *((wrapper, "simulation") for wrapper in WRAPPERS), + (Simulation, "holder"), +) +CORE_WRITE_ROUTES = (*WRAPPER_WRITE_ROUTES, (make_core_simulation, "simulation")) + + +def write_public_input(simulation, name, input_period, values, route): + if route == "holder": + return simulation.get_holder(name).set_input( + period(input_period), np.asarray(values), simulation.branch_name + ) + return simulation.set_input(name, input_period, values) + + +def core_builder_wrapper(*, situation, tax_benefit_system): + """Use raw Core with the same constructor shape as the public wrappers.""" + return SimulationBuilder().build_from_dict(tax_benefit_system, situation) + + +PERIOD_WRAPPERS = (*WRAPPERS, core_builder_wrapper) + + def amount(simulation, name, period): return np.asarray(simulation.calculate(name, period))[0] @@ -181,10 +208,10 @@ def test_core_builder_migrates_initial_public_inputs_and_preserves_precedence(ro assert amount(simulation, f"{disability}_reported", 2024) == 666 -@pytest.mark.parametrize("wrapper", (*WRAPPERS, make_core_simulation)) +@pytest.mark.parametrize("wrapper,write_route", CORE_WRITE_ROUTES) @pytest.mark.parametrize("reported_input", (None, 0, 6_000)) def test_later_public_writes_refresh_automatic_sources_and_keep_explicit_twins( - wrapper, reported_input + wrapper, write_route, reported_input ): public = SOCIAL_SECURITY_COMPONENTS[0] reported = f"{public}_reported" @@ -197,7 +224,7 @@ def test_later_public_writes_refresh_automatic_sources_and_keep_explicit_twins( # A calculated twin, even a cached zero, is not an explicitly supplied twin. assert amount(simulation, reported, 2024) == (reported_input or 0) simulation.set_input(FACTOR, 2024, [1.08]) - simulation.set_input(public, 2024, [12_000]) + write_public_input(simulation, public, 2024, [12_000], write_route) expected_source = 12_000 if reported_input is None else reported_input assert amount(simulation, public, 2024) == 12_000 assert amount(simulation, reported, 2024) == expected_source @@ -207,7 +234,7 @@ def test_later_public_writes_refresh_automatic_sources_and_keep_explicit_twins( assert amount(simulation, public, 2026) == pytest.approx( expected_source * 13_522.16 / 12_000, abs=0.01 ) - simulation.set_input(public, 2024, [24_000]) + write_public_input(simulation, public, 2024, [24_000], write_route) expected_source = 24_000 if reported_input is None else reported_input assert amount(simulation, reported, 2024) == expected_source drop_inherited_values(simulation) @@ -288,8 +315,10 @@ def test_later_monthly_public_write_seeds_the_stored_rolling_year(wrapper): assert amount(simulation, public, 2026) > 12_000 -@pytest.mark.parametrize("wrapper", WRAPPERS) -def test_clone_and_branch_later_public_writes_keep_input_lists_private(wrapper): +@pytest.mark.parametrize("wrapper,write_route", WRAPPER_WRITE_ROUTES) +def test_clone_and_branch_later_public_writes_keep_input_lists_private( + wrapper, write_route +): public = SOCIAL_SECURITY_COMPONENTS[0] reported = f"{public}_reported" simulation = make_simulation(wrapper, {public: {2024: 12_000}}) @@ -300,7 +329,7 @@ def test_clone_and_branch_later_public_writes_keep_input_lists_private(wrapper): simulation.clone(clone_tax_benefit_system=False), simulation.get_branch("new_public_override"), ): - copy.set_input(public, 2025, [24_000]) + write_public_input(copy, public, 2025, [24_000], write_route) drop_inherited_values(copy) assert amount(copy, public, 2025) == 24_000 assert amount(copy, reported, 2025) == 24_000 @@ -361,7 +390,10 @@ def test_neutral_reform_replay_keeps_automatic_source_origin_private(wrapper): assert amount(simulation, public, 2026) == parent_later -def test_dataset_constructor_migrates_stored_original_columns(): +@pytest.mark.parametrize( + "wrapper,factor", ((Microsimulation, 1.1), (CoreSimulation, 1.08)) +) +def test_dataset_constructor_migrates_stored_original_columns(wrapper, factor): public = SOCIAL_SECURITY_COMPONENTS[0] frame = pd.DataFrame( { @@ -391,10 +423,17 @@ def test_dataset_constructor_migrates_stored_original_columns(): f"{public}_reported__2026": [13_000], } ) - simulation = Microsimulation(dataset=frame, tax_benefit_system=system) - simulation.set_input(FACTOR, 2024, [1.1]) - assert amount(simulation, public, 2024) == pytest.approx(13_200) + simulation = wrapper( + dataset=frame, + tax_benefit_system=system, + default_input_period=2024, + default_calculation_period=2024, + ) + simulation.set_input(FACTOR, 2024, [factor]) + simulation.set_input(FACTOR, 2026, [factor]) + assert amount(simulation, public, 2024) == pytest.approx(12_000 * factor) assert amount(simulation, public, 2026) == 14_000 + assert amount(simulation, f"{public}_reported", 2026) == 13_000 assert input_periods(simulation, public) == {"2026"} assert input_periods(simulation, f"{public}_reported") == {"2024", "2026"} # Loader-time public writes migrate; later writes remain final outputs. @@ -538,15 +577,15 @@ def test_legacy_extraction_keeps_later_override_and_matching_source( ) -@pytest.mark.parametrize("wrapper", WRAPPERS) +@pytest.mark.parametrize("wrapper,write_route", WRAPPER_WRITE_ROUTES) @pytest.mark.parametrize("snapshot", ("legacy_extraction", "complete_export")) def test_source_round_trip_retains_automatic_origin_for_subsequent_writes( - wrapper, snapshot + wrapper, write_route, snapshot ): public = SOCIAL_SECURITY_COMPONENTS[0] reported = f"{public}_reported" simulation = make_simulation(wrapper, {public: {2024: 12_000}}) - simulation.set_input(public, 2024, [24_000]) + write_public_input(simulation, public, 2024, [24_000], write_route) # A cached intervening source is not exported or treated as supplied. amount(simulation, reported, 2025) if snapshot == "legacy_extraction": @@ -562,7 +601,7 @@ def test_source_round_trip_retains_automatic_origin_for_subsequent_writes( assert payload["overrides"][public] == {"2024": [24_000]} assert person[reported] == {"2024": 24_000} assert set(person).issubset(system.variables) - reloaded.set_input(public, 2024, [36_000]) + write_public_input(reloaded, public, 2024, [36_000], write_route) assert amount(reloaded, public, 2024) == 36_000 assert amount(reloaded, reported, 2024) == 36_000 assert amount(reloaded, public, 2026) == pytest.approx(40_566.48, abs=0.02) @@ -700,7 +739,7 @@ def test_parameter_reform_registration_order_preserves_stable_public_formula(): ) -@pytest.mark.parametrize("wrapper", WRAPPERS) +@pytest.mark.parametrize("wrapper", PERIOD_WRAPPERS) def test_exact_multi_year_sources_keep_keys_uniform_factors_and_overrides(wrapper): public = SOCIAL_SECURITY_COMPONENTS[0] reported = f"{public}_reported" @@ -739,29 +778,33 @@ def test_exact_multi_year_sources_keep_keys_uniform_factors_and_overrides(wrappe assert input_periods(override_reload, "household_weight") == {multi_year} +@pytest.mark.parametrize("wrapper", PERIOD_WRAPPERS) def test_exact_multi_year_both_names_and_neutralization_keep_core_precedence( - monkeypatch, + monkeypatch, wrapper ): public = SOCIAL_SECURITY_COMPONENTS[0] reported = f"{public}_reported" multi_year = "year:2024:2" simulation = make_simulation( - inputs={public: {multi_year: 123}, reported: {multi_year: 24_000}} + wrapper, {public: {multi_year: 123}, reported: {multi_year: 24_000}} ) simulation.set_input(FACTOR, 2024, [1.1]) simulation.set_input(FACTOR, 2025, [0.8]) assert amount(simulation, public, multi_year) == 123 assert amount(simulation, reported, multi_year) == 24_000 - neutralized = make_simulation(inputs={reported: {multi_year: 24_000}}) + neutralized = make_simulation(wrapper, {reported: {multi_year: 24_000}}) with monkeypatch.context() as context: context.setattr(system.variables[public], "is_neutralized", True) assert amount(neutralized, public, multi_year) == 0 -def test_exact_multi_year_source_rejects_differing_annual_factors_without_allocation(): +@pytest.mark.parametrize("wrapper", PERIOD_WRAPPERS) +def test_exact_multi_year_source_rejects_differing_annual_factors_without_allocation( + wrapper, +): public = SOCIAL_SECURITY_COMPONENTS[0] multi_year = "year:2024:2" - simulation = make_simulation(inputs={public: {multi_year: 24_000}}) + simulation = make_simulation(wrapper, {public: {multi_year: 24_000}}) simulation.set_input(FACTOR, 2024, [1.1]) simulation.set_input(FACTOR, 2025, [0.8]) with pytest.raises(ValueError, match="separately supplied annual source amounts"): @@ -790,12 +833,15 @@ def test_containing_period_migration_preserves_sources_and_inner_final_overrides assert amount(simulation, reported, "year:2024:10") == 100_000 -def test_multi_year_selected_base_input_and_neutralization_take_precedence(monkeypatch): +@pytest.mark.parametrize("wrapper", PERIOD_WRAPPERS) +def test_multi_year_selected_base_input_and_neutralization_take_precedence( + monkeypatch, wrapper +): public = SOCIAL_SECURITY_COMPONENTS[0] reported, base = f"{public}_reported", f"{public}_base" multi_year = "year:2024:2" simulation = make_simulation( - inputs={reported: {multi_year: 24_000}, base: {multi_year: 6_000}} + wrapper, {reported: {multi_year: 24_000}, base: {multi_year: 6_000}} ) for year in (2024, 2025): simulation.set_input(FACTOR, year, [1.1]) @@ -805,10 +851,10 @@ def test_multi_year_selected_base_input_and_neutralization_take_precedence(monke ) assert amount(reloaded, public, multi_year) == pytest.approx(6_600) - base_only = make_simulation(inputs={base: {multi_year: 6_000}}) + base_only = make_simulation(wrapper, {base: {multi_year: 6_000}}) assert amount(base_only, public, multi_year) == 6_000 neutralized = make_simulation( - inputs={reported: {multi_year: 24_000}, base: {multi_year: 6_000}} + wrapper, {reported: {multi_year: 24_000}, base: {multi_year: 6_000}} ) neutralized.set_input(FACTOR, 2024, [1.1]) neutralized.set_input(FACTOR, 2025, [0.8]) diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 4593ccbdccd..7648b9675b8 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -111,18 +111,22 @@ def install_social_security_input_compatibility(): """Extend Core's ingress routes only for systems opting into this policy. Core's entity builder writes holders directly, and its variable-level - ``set_input`` handler runs only for mismatched period units. A builder - adapter therefore migrates initial sources, while a simulation adapter - preserves subsequent public overrides, including on raw Core simulations. + ``set_input`` handler runs only for mismatched period units. Builder + and dataset adapters migrate initial sources. Holder-level interception + preserves subsequent public overrides from every supported input route: + simulation setters, direct holder writes, and custom period helpers. """ from policyengine_core.simulations import Simulation as CoreSimulation from policyengine_core.simulations.simulation_builder import SimulationBuilder + from policyengine_core.holders import Holder - original_set_input = CoreSimulation.set_input + original_set_input = Holder.set_input if not getattr(original_set_input, "_social_security_compatibility", False): @wraps(original_set_input) - def set_input(simulation, variable_name, period, value): + def set_input(holder, period, value, branch_name="default"): + simulation = holder.simulation + variable_name = holder.variable.name track_public_write = ( variable_name in SOCIAL_SECURITY_COMPONENTS and getattr( @@ -134,7 +138,7 @@ def set_input(simulation, variable_name, period, value): and not _initial_social_security_ingress.get() ) if not track_public_write: - return original_set_input(simulation, variable_name, period, value) + return original_set_input(holder, period, value, branch_name) before = { supplied_period: _recorded_input_location( @@ -142,7 +146,7 @@ def set_input(simulation, variable_name, period, value): ) for supplied_period in simulation._get_set_input_periods(variable_name) } - result = original_set_input(simulation, variable_name, period, value) + result = original_set_input(holder, period, value, branch_name) reported_name = f"{variable_name}_reported" reported_periods = set(simulation._get_set_input_periods(reported_name)) # Period helpers can store a different period from the one passed @@ -163,10 +167,10 @@ def set_input(simulation, variable_name, period, value): ) if supplied_value is not None: original_set_input( - simulation, - reported_name, + simulation.get_holder(reported_name), supplied_period, supplied_value.copy(), + branch_name, ) _mark_automatic_reported_input( simulation, reported_name, supplied_period @@ -182,7 +186,32 @@ def set_input(simulation, variable_name, period, value): return result set_input._social_security_compatibility = True - CoreSimulation.set_input = set_input + Holder.set_input = set_input + + original_dataset_build = CoreSimulation.build_from_dataset + if not getattr(original_dataset_build, "_social_security_compatibility", False): + + @wraps(original_dataset_build) + def build_from_dataset(simulation): + if not getattr( + simulation.tax_benefit_system, + "_social_security_input_compatibility", + False, + ): + return original_dataset_build(simulation) + token = _initial_social_security_ingress.set(True) + try: + result = original_dataset_build(simulation) + finally: + _initial_social_security_ingress.reset(token) + # Raw Core simulations use this loader without either country + # constructor or an entity/variable builder. Migrate before Core + # creates its reform baseline, exactly as for situation ingress. + migrate_social_security_inputs(simulation) + return result + + build_from_dataset._social_security_compatibility = True + CoreSimulation.build_from_dataset = build_from_dataset original_extract = CoreSimulation.extract_person if not getattr(original_extract, "_social_security_compatibility", False): @@ -331,6 +360,12 @@ def build(builder, tax_benefit_system, input_dict, simulation=None): if not getattr(original_build, "_social_security_compatibility", False): setattr(SimulationBuilder, method_name, wrap_builder(original_build)) + from policyengine_us.tools.social_security_periods import ( + install_social_security_period_compatibility, + ) + + install_social_security_period_compatibility() + def migrate_social_security_inputs(simulation, *, include_baseline=True): """Move legacy component inputs, period by period, into pure sources. diff --git a/policyengine_us/tools/social_security_periods.py b/policyengine_us/tools/social_security_periods.py index ba3b60eda6d..451e44bdbb8 100644 --- a/policyengine_us/tools/social_security_periods.py +++ b/policyengine_us/tools/social_security_periods.py @@ -1,8 +1,9 @@ -"""Preserve exact multi-year benefit sources in PE-US wrapper simulations.""" +"""Preserve exact multi-year benefit sources through every Core route.""" + +from functools import wraps import numpy as np from policyengine_core.periods import YEAR -from policyengine_core.periods import period as to_period from policyengine_core.simulations import Simulation as CoreSimulation from policyengine_us.tools.social_security_benefit_levers import ( @@ -12,115 +13,133 @@ from policyengine_us.tools.social_security_inputs import SOCIAL_SECURITY_COMPONENTS -class SocialSecurityPeriodMixin: - def _calculate(self, variable_name, period=None): - if ( - variable_name not in SOCIAL_SECURITY_COMPONENTS - or period is None - or period.unit != YEAR - or period.size == 1 - ): - return super()._calculate(variable_name, period) - variable = self.tax_benefit_system.get_variable(variable_name) - # Core checks neutralization before exact public inputs and caches. - if variable.is_neutralized: - return super()._calculate(variable_name, period) - try: - abolished = self.tax_benefit_system.parameters(period).gov.abolitions[ - variable_name - ] - except Exception: - abolished = False - if abolished: - return super()._calculate(variable_name, period) - holder = self.get_holder(variable_name) - if holder.get_array(period, self.branch_name) is not None: - return super()._calculate(variable_name, period) - reported_name = f"{variable_name}_reported" - base_name = f"{variable_name}_base" - base_periods = self._get_set_input_periods(base_name) - has_exact_base = period in base_periods - if not has_exact_base and period not in self._get_set_input_periods( - reported_name +def install_social_security_period_compatibility(): + """Apply the compound-period rule to opted-in raw and wrapper simulations.""" + original_calculate = CoreSimulation._calculate + if getattr(original_calculate, "_social_security_period_compatibility", False): + return + + @wraps(original_calculate) + def calculate(simulation, variable_name, period=None): + if not getattr( + simulation.tax_benefit_system, + "_social_security_input_compatibility", + False, ): - return super()._calculate(variable_name, period) + return original_calculate(simulation, variable_name, period) + return _calculate_social_security_period( + simulation, variable_name, period, original_calculate + ) + + calculate._social_security_period_compatibility = True + CoreSimulation._calculate = calculate + + +def _calculate_social_security_period(self, variable_name, period, original_calculate): + if ( + variable_name not in SOCIAL_SECURITY_COMPONENTS + or period is None + or period.unit != YEAR + or period.size == 1 + ): + return original_calculate(self, variable_name, period) + variable = self.tax_benefit_system.get_variable(variable_name) + # Core checks neutralization before exact public inputs and caches. + if variable.is_neutralized: + return original_calculate(self, variable_name, period) + try: + abolished = self.tax_benefit_system.parameters(period).gov.abolitions[ + variable_name + ] + except Exception: + abolished = False + if abolished: + return original_calculate(self, variable_name, period) + holder = self.get_holder(variable_name) + if holder.get_array(period, self.branch_name) is not None: + return original_calculate(self, variable_name, period) + reported_name = f"{variable_name}_reported" + base_name = f"{variable_name}_base" + base_periods = self._get_set_input_periods(base_name) + has_exact_base = period in base_periods + if not has_exact_base and period not in self._get_set_input_periods(reported_name): + return original_calculate(self, variable_name, period) - input_state = self._calculation_start() - base_variable = self.tax_benefit_system.get_variable(base_name) + input_state = self._calculation_start() + base_variable = self.tax_benefit_system.get_variable(base_name) + try: + base_abolished = self.tax_benefit_system.parameters(period).gov.abolitions[ + base_name + ] + except Exception: + base_abolished = False + if has_exact_base or base_variable.is_neutralized or base_abolished: + # Let Core preserve exact base-input and neutralization precedence. + amount = CoreSimulation.calculate(self, base_name, period) + else: + if any( + known.start <= period.stop and known.stop >= period.start + for known in base_periods + ): + raise ValueError( + "A multi-year reported input with base overrides for other " + "periods requires separately supplied annual source amounts" + ) + amount = CoreSimulation.calculate(self, reported_name, period) + nonzero = amount != 0 + # Core accepts a stored multi-year YEAR input before checking period + # consistency. Preserve that exact total: set_input_divide_by_period would + # invent annual inputs and change both annual results and future uprating. + # A uniform multiplier needs no allocation; differing annual multipliers + # need separately supplied annual sources. Core supplies the subperiods. + common_factor = None + for year in period.get_subperiods(YEAR): + source = CoreSimulation.calculate(self, f"{variable_name}_source", year) + if np.any(source != SocialSecurityBenefitSource.reported): + raise NotImplementedError( + "Multi-year Social Security sources require the reported route; " + "calculated source modes are reserved" + ) + factor = np.ones_like(amount) + for name in ( + f"{variable_name}_joint_cola_factor", + f"{variable_name}_pia_ratio_factor", + "social_security_across_the_board_factor", + f"{variable_name}_payable_factor", + ): + factor = factor * CoreSimulation.calculate(self, name, year) try: - base_abolished = self.tax_benefit_system.parameters(period).gov.abolitions[ - base_name - ] + annual_base_abolished = self.tax_benefit_system.parameters( + year + ).gov.abolitions[base_name] except Exception: - base_abolished = False - if has_exact_base or base_variable.is_neutralized or base_abolished: - # Let Core preserve exact base-input and neutralization precedence. - amount = CoreSimulation.calculate(self, base_name, period) - else: - if any( - known.start <= period.stop and known.stop >= period.start - for known in base_periods - ): - raise ValueError( - "A multi-year reported input with base overrides for other " - "periods requires separately supplied annual source amounts" - ) - amount = CoreSimulation.calculate(self, reported_name, period) - nonzero = amount != 0 - # Core accepts a stored multi-year YEAR input before checking period - # consistency. A uniform multiplier needs no annual allocation, while - # differing annual multipliers need separately supplied annual sources. - common_factor = None - for year_number in range(period.start.year, period.stop.year + 1): - year = to_period(year_number) - source = CoreSimulation.calculate(self, f"{variable_name}_source", year) - if np.any(source != SocialSecurityBenefitSource.reported): - raise NotImplementedError( - "Multi-year Social Security sources require the reported route; " - "calculated source modes are reserved" - ) - factor = np.ones_like(amount) - for name in ( - f"{variable_name}_joint_cola_factor", - f"{variable_name}_pia_ratio_factor", - "social_security_across_the_board_factor", - f"{variable_name}_payable_factor", - ): - factor = factor * CoreSimulation.calculate(self, name, year) - try: - annual_base_abolished = self.tax_benefit_system.parameters( - year - ).gov.abolitions[base_name] - except Exception: - annual_base_abolished = False - if annual_base_abolished: - factor = np.zeros_like(factor) - if ( - variable_name == "social_security_dependents" - and self.tax_benefit_system.parameters( - year - ).gov.contrib.ssa.benefit_levers.payable.oasi_only - and np.any( - CoreSimulation.calculate( - self, "social_security_payable_factor", year - ) - != 1 - ) - ): - fund = CoreSimulation.calculate( - self, "social_security_dependents_fund", year - ) - if np.any((amount > 0) & (fund == SocialSecurityFund.unspecified)): - raise ValueError( - "OASI-only payable cuts require the dependent benefit fund" - ) - if common_factor is None: - common_factor = factor - elif not np.array_equal(common_factor[nonzero], factor[nonzero]): + annual_base_abolished = False + if annual_base_abolished: + factor = np.zeros_like(factor) + if ( + variable_name == "social_security_dependents" + and self.tax_benefit_system.parameters( + year + ).gov.contrib.ssa.benefit_levers.payable.oasi_only + and np.any( + CoreSimulation.calculate(self, "social_security_payable_factor", year) + != 1 + ) + ): + fund = CoreSimulation.calculate( + self, "social_security_dependents_fund", year + ) + if np.any((amount > 0) & (fund == SocialSecurityFund.unspecified)): raise ValueError( - "Social Security levers with differing annual factors on a " - "multi-year reported input require separately supplied " - "annual source amounts" + "OASI-only payable cuts require the dependent benefit fund" ) - result = self._cast_formula_result(amount * common_factor, variable) - return self._cache_result(holder, result, period, input_state) + if common_factor is None: + common_factor = factor + elif not np.array_equal(common_factor[nonzero], factor[nonzero]): + raise ValueError( + "Social Security levers with differing annual factors on a " + "multi-year reported input require separately supplied " + "annual source amounts" + ) + result = self._cast_formula_result(amount * common_factor, variable) + return self._cache_result(holder, result, period, input_state) From 02b741510c273043666f2c1d12aa38e43c79d20c Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 24 Oct 2026 09:11:18 -0400 Subject: [PATCH 27/76] Cover branch slots and direct Social Security input helpers --- .../test_social_security_inputs.py | 54 +++++- .../tools/social_security_inputs.py | 169 ++++++++++++------ 2 files changed, 166 insertions(+), 57 deletions(-) diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index b787e462cc6..fc5eb8352fa 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -10,6 +10,10 @@ import numpy as np import pandas as pd import pytest +from policyengine_core.holders.helpers import ( + set_input_dispatch_by_period, + set_input_divide_by_period, +) from policyengine_core.periods import period from policyengine_core.simulations import ( Simulation as CoreSimulation, @@ -79,10 +83,24 @@ def make_core_simulation(inputs=None, route="dict_entities"): *((wrapper, "simulation") for wrapper in WRAPPERS), (Simulation, "holder"), ) -CORE_WRITE_ROUTES = (*WRAPPER_WRITE_ROUTES, (make_core_simulation, "simulation")) +CORE_WRITE_ROUTES = ( + *WRAPPER_WRITE_ROUTES, + (make_core_simulation, "simulation"), + (make_core_simulation, "dispatch"), + (make_core_simulation, "divide"), +) def write_public_input(simulation, name, input_period, values, route): + if route in ("dispatch", "divide"): + helper = ( + set_input_dispatch_by_period + if route == "dispatch" + else set_input_divide_by_period + ) + return helper( + simulation.get_holder(name), period(input_period), np.asarray(values) + ) if route == "holder": return simulation.get_holder(name).set_input( period=period(input_period), @@ -236,8 +254,11 @@ def test_later_public_writes_refresh_automatic_sources_and_keep_explicit_twins( assert amount(simulation, public, 2026) == pytest.approx( expected_source * 13_522.16 / 12_000, abs=0.01 ) - write_public_input(simulation, public, 2024, [24_000], write_route) - expected_source = 24_000 if reported_input is None else reported_input + # Core's period-splitting helpers preserve an already supplied annual + # input, and divide rejects a conflicting repeated total. + replacement = 12_000 if write_route in ("dispatch", "divide") else 24_000 + write_public_input(simulation, public, 2024, [replacement], write_route) + expected_source = replacement if reported_input is None else reported_input assert amount(simulation, reported, 2024) == expected_source drop_inherited_values(simulation) assert amount(simulation, public, 2026) == pytest.approx( @@ -344,6 +365,33 @@ def test_clone_and_branch_later_public_writes_keep_input_lists_private( assert input_periods(simulation, reported) == {"2024"} +def test_holder_write_before_branch_creation_stays_in_its_named_branch(): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + simulation = make_simulation() + simulation.get_holder(public).set_input( + period=period(2024), array=np.array([12_000]), branch_name="later" + ) + branch = simulation.get_branch("later") + assert amount(branch, public, 2024) == 12_000 + assert amount(branch, reported, 2024) == 12_000 + assert amount(branch, public, 2026) == pytest.approx(13_522.16, abs=0.01) + for year in (2024, 2026): + assert amount(simulation, public, year) == 0 + assert amount(simulation, reported, year) == 0 + + reloaded = Simulation(situation=branch.extract_person(0), tax_benefit_system=system) + assert amount(reloaded, public, 2024) == 12_000 + assert amount(reloaded, public, 2026) == pytest.approx(13_522.16, abs=0.01) + clone = branch.clone(clone_tax_benefit_system=False) + write_public_input(clone, public, 2024, [24_000], "holder") + drop_inherited_values(clone) + assert amount(clone, public, 2026) == pytest.approx(27_044.32, abs=0.01) + assert amount(branch, public, 2026) == pytest.approx(13_522.16, abs=0.01) + assert amount(reloaded, public, 2026) == pytest.approx(13_522.16, abs=0.01) + assert amount(simulation, public, 2026) == 0 + + @pytest.mark.parametrize("wrapper", WRAPPERS) def test_branch_explicit_twin_overrides_inherited_automatic_source_privately(wrapper): public = SOCIAL_SECURITY_COMPONENTS[0] diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 9db278a6074..e57fead21e4 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -9,6 +9,7 @@ from contextvars import ContextVar from functools import wraps +from policyengine_core.holders.helpers import get_stored_input from policyengine_core.periods import period as to_period from policyengine_core.simulations.simulation import NpEncoder @@ -54,8 +55,10 @@ def _recorded_input_location(simulation, name, period, *, branch_names=None): return None -def _is_automatic_reported_input(simulation, name, period): - location = _recorded_input_location(simulation, name, period) +def _is_automatic_reported_input(simulation, name, period, *, branch_names=None): + location = _recorded_input_location( + simulation, name, period, branch_names=branch_names + ) return ( location is not None and getattr(simulation, _AUTOMATIC_INPUTS, {}).get(location[0]) == location[1] @@ -120,73 +123,131 @@ def install_social_security_input_compatibility(): from policyengine_core.simulations.simulation_builder import SimulationBuilder from policyengine_core.holders import Holder + def tracks_public_input(holder): + simulation = getattr(holder, "simulation", None) + return ( + holder.variable.name in SOCIAL_SECURITY_COMPONENTS + and getattr( + getattr(simulation, "tax_benefit_system", None), + "_social_security_input_compatibility", + False, + ) + and hasattr(simulation, "input_variables") + and not _initial_social_security_ingress.get() + ) + + def input_locations(holder, branch_name): + return { + period: location + for stored_branch, period in holder.get_known_branch_periods() + if stored_branch == branch_name + and ( + location := _recorded_input_location( + holder.simulation, + holder.variable.name, + period, + branch_names=(branch_name,), + ) + ) + is not None + } + + def synchronize_public_input(holder, period, branch_name): + simulation = holder.simulation + variable_name = holder.variable.name + period = holder._storage_period(period) + if ( + _recorded_input_location( + simulation, variable_name, period, branch_names=(branch_name,) + ) + is None + ): + return + reported_name = f"{variable_name}_reported" + reported_holder = simulation.get_holder(reported_name) + readable_branches = reported_holder._readable_branches(branch_name) + reported_location = _recorded_input_location( + simulation, reported_name, period, branch_names=readable_branches + ) + if reported_location is None or _is_automatic_reported_input( + simulation, reported_name, period, branch_names=readable_branches + ): + supplied_value = get_stored_input(holder, period, branch_name) + if supplied_value is not None: + reported_holder.set_input(period, supplied_value.copy(), branch_name) + _mark_automatic_reported_input( + simulation, reported_name, period, branch_name=branch_name + ) + # Core clones share this legacy list. Replace it privately so a + # later override is extracted without changing the source clone. + simulation.input_variables = list( + dict.fromkeys([*simulation.input_variables, variable_name, reported_name]) + ) + original_set_input = Holder.set_input if not getattr(original_set_input, "_social_security_compatibility", False): @wraps(original_set_input) def set_input(holder, period, array, branch_name="default"): - simulation = holder.simulation - variable_name = holder.variable.name - track_public_write = ( - variable_name in SOCIAL_SECURITY_COMPONENTS - and getattr( - simulation.tax_benefit_system, - "_social_security_input_compatibility", - False, - ) - and hasattr(simulation, "input_variables") - and not _initial_social_security_ingress.get() - ) - if not track_public_write: + if not tracks_public_input(holder): return original_set_input(holder, period, array, branch_name) - before = { - supplied_period: _recorded_input_location( - simulation, variable_name, supplied_period - ) - for supplied_period in simulation._get_set_input_periods(variable_name) - } + before = input_locations(holder, branch_name) result = original_set_input(holder, period, array, branch_name) - reported_name = f"{variable_name}_reported" - reported_periods = set(simulation._get_set_input_periods(reported_name)) # Period helpers can store a different period from the one passed # in: a monthly input to a yearly variable stores a rolling year. - # Inspect actual stores rather than guessing those period aliases. - for supplied_period in simulation._get_set_input_periods(variable_name): - if _recorded_input_location( - simulation, variable_name, supplied_period - ) == before.get(supplied_period) or ( - supplied_period in reported_periods - and not _is_automatic_reported_input( - simulation, reported_name, supplied_period - ) - ): - continue - supplied_value = simulation._get_recorded_input_array( - variable_name, supplied_period - ) - if supplied_value is not None: - original_set_input( - simulation.get_holder(reported_name), - supplied_period, - supplied_value.copy(), - branch_name, - ) - _mark_automatic_reported_input( - simulation, reported_name, supplied_period - ) + # Inspect actual inputs under the written branch, including a + # branch slot this simulation will read only after cloning. + for supplied_period, location in input_locations( + holder, branch_name + ).items(): + if location != before.get(supplied_period): + synchronize_public_input(holder, supplied_period, branch_name) + return result - # Core clones share this legacy list. Replace it privately so a - # later override is extracted without changing the source clone. - simulation.input_variables = list( - dict.fromkeys( - [*simulation.input_variables, variable_name, reported_name] + set_input._social_security_compatibility = True + Holder.set_input = set_input + + original_store = Holder._set + if not getattr(original_store, "_social_security_compatibility", False): + + @wraps(original_store) + def store_input( + holder, + period, + value, + branch_name="default", + validate_nan=False, + derived=False, + is_input=None, + sequence_number=None, + ): + # Core's direct period helpers and documented custom input stores + # can call _set(is_input=True) without entering Holder.set_input. + # Normal setter calls synchronize once, after their handler ends. + track_direct_input = ( + is_input + and not getattr( + getattr(holder, "simulation", None), "_user_input_contexts", None ) + and tracks_public_input(holder) + ) + result = original_store( + holder, + period, + value, + branch_name, + validate_nan=validate_nan, + derived=derived, + is_input=is_input, + sequence_number=sequence_number, ) + if track_direct_input: + synchronize_public_input(holder, period, branch_name) return result - set_input._social_security_compatibility = True - Holder.set_input = set_input + store_input._social_security_compatibility = True + Holder._set = store_input original_dataset_build = CoreSimulation.build_from_dataset if not getattr(original_dataset_build, "_social_security_compatibility", False): From 5f4b175df4b3bbc33a576141e15855e961102ecb Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 24 Oct 2026 09:28:20 -0400 Subject: [PATCH 28/76] Supply exact-period weights in compound Microsimulation fixtures --- .../tests/microsimulation/test_social_security_inputs.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index fc5eb8352fa..604694a4e9d 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -838,11 +838,13 @@ def test_exact_multi_year_both_names_and_neutralization_keep_core_precedence( simulation = make_simulation( wrapper, {public: {multi_year: 123}, reported: {multi_year: 24_000}} ) + simulation.set_input("household_weight", multi_year, [1]) simulation.set_input(FACTOR, 2024, [1.1]) simulation.set_input(FACTOR, 2025, [0.8]) assert amount(simulation, public, multi_year) == 123 assert amount(simulation, reported, multi_year) == 24_000 neutralized = make_simulation(wrapper, {reported: {multi_year: 24_000}}) + neutralized.set_input("household_weight", multi_year, [1]) with monkeypatch.context() as context: context.setattr(system.variables[public], "is_neutralized", True) assert amount(neutralized, public, multi_year) == 0 @@ -893,6 +895,7 @@ def test_multi_year_selected_base_input_and_neutralization_take_precedence( simulation = make_simulation( wrapper, {reported: {multi_year: 24_000}, base: {multi_year: 6_000}} ) + simulation.set_input("household_weight", multi_year, [1]) for year in (2024, 2025): simulation.set_input(FACTOR, year, [1.1]) assert amount(simulation, public, multi_year) == pytest.approx(6_600) @@ -902,10 +905,12 @@ def test_multi_year_selected_base_input_and_neutralization_take_precedence( assert amount(reloaded, public, multi_year) == pytest.approx(6_600) base_only = make_simulation(wrapper, {base: {multi_year: 6_000}}) + base_only.set_input("household_weight", multi_year, [1]) assert amount(base_only, public, multi_year) == 6_000 neutralized = make_simulation( wrapper, {reported: {multi_year: 24_000}, base: {multi_year: 6_000}} ) + neutralized.set_input("household_weight", multi_year, [1]) neutralized.set_input(FACTOR, 2024, [1.1]) neutralized.set_input(FACTOR, 2025, [0.8]) with monkeypatch.context() as context: From 01b60b27f867deeabdfd28c81b283ada41b3a8d0 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 24 Oct 2026 09:57:15 -0400 Subject: [PATCH 29/76] Keep calendar-year factor checks for rolling compound inputs --- .../test_social_security_inputs.py | 16 ++++++++++++---- policyengine_us/tools/social_security_periods.py | 7 ++++++- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index 604694a4e9d..1ded519fbf4 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -851,14 +851,22 @@ def test_exact_multi_year_both_names_and_neutralization_keep_core_precedence( @pytest.mark.parametrize("wrapper", PERIOD_WRAPPERS) +@pytest.mark.parametrize( + "multi_year,annual_factors", + ( + ("year:2024:2", {2024: 1.1, 2025: 0.8}), + ("year:2024-07:2", {2024: 1.1, 2025: 1.1, 2026: 0.8}), + ), +) def test_exact_multi_year_source_rejects_differing_annual_factors_without_allocation( - wrapper, + wrapper, multi_year, annual_factors ): public = SOCIAL_SECURITY_COMPONENTS[0] - multi_year = "year:2024:2" simulation = make_simulation(wrapper, {public: {multi_year: 24_000}}) - simulation.set_input(FACTOR, 2024, [1.1]) - simulation.set_input(FACTOR, 2025, [0.8]) + simulation.set_input("household_weight", multi_year, [1]) + assert amount(simulation, public, multi_year) == 24_000 + for year, factor in annual_factors.items(): + simulation.set_input(FACTOR, year, [factor]) with pytest.raises(ValueError, match="separately supplied annual source amounts"): simulation.calculate(public, multi_year) assert input_periods(simulation, f"{public}_reported") == {multi_year} diff --git a/policyengine_us/tools/social_security_periods.py b/policyengine_us/tools/social_security_periods.py index 451e44bdbb8..ed8a59cc93b 100644 --- a/policyengine_us/tools/social_security_periods.py +++ b/policyengine_us/tools/social_security_periods.py @@ -93,7 +93,12 @@ def _calculate_social_security_period(self, variable_name, period, original_calc # A uniform multiplier needs no allocation; differing annual multipliers # need separately supplied annual sources. Core supplies the subperiods. common_factor = None - for year in period.get_subperiods(YEAR): + # A rolling compound can overlap one more calendar year than its size. + # Check every calendar year touched, including the final partial year. + calendar_span = period.this_year.start.period( + YEAR, size=period.stop.year - period.start.year + 1 + ) + for year in calendar_span.get_subperiods(YEAR): source = CoreSimulation.calculate(self, f"{variable_name}_source", year) if np.any(source != SocialSecurityBenefitSource.reported): raise NotImplementedError( From f164c2190f937e1354c0a998c90bcb7f7e3e6215 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 24 Oct 2026 10:02:35 -0400 Subject: [PATCH 30/76] Clear compound test caches before rechecking annual factors --- .../tests/microsimulation/test_social_security_inputs.py | 1 + 1 file changed, 1 insertion(+) diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index 1ded519fbf4..fc0df93e547 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -867,6 +867,7 @@ def test_exact_multi_year_source_rejects_differing_annual_factors_without_alloca assert amount(simulation, public, multi_year) == 24_000 for year, factor in annual_factors.items(): simulation.set_input(FACTOR, year, [factor]) + drop_inherited_values(simulation) with pytest.raises(ValueError, match="separately supplied annual source amounts"): simulation.calculate(public, multi_year) assert input_periods(simulation, f"{public}_reported") == {multi_year} From 1f28500ad308d550c03cfcb35d21b4615dac5db6 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 01:11:53 -0400 Subject: [PATCH 31/76] Preserve Social Security sources across Core input routes --- changelog.d/ss-benefit-levers.added.md | 2 +- .../test_social_security_current_law.py | 153 +++++++++- .../test_social_security_inputs.py | 186 +++++++++++++ .../tools/social_security_inputs.py | 261 ++++++++++++------ 4 files changed, 520 insertions(+), 82 deletions(-) diff --git a/changelog.d/ss-benefit-levers.added.md b/changelog.d/ss-benefit-levers.added.md index cec2a318738..21001b7f1e4 100644 --- a/changelog.d/ss-benefit-levers.added.md +++ b/changelog.d/ss-benefit-levers.added.md @@ -1 +1 @@ -Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs migrate on every ingress route; later public overrides refresh automatic sources, preserve explicit reported inputs, and survive extraction. +Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs migrate on every ingress route; later public overrides refresh automatic sources, preserve explicit reported inputs, and survive extraction. Core dump/restore and subsampling preserve source provenance; direct input-context stores and legacy nonderived cache seeds retain future uprating. diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 223884644e7..136143dbfa3 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -139,6 +139,12 @@ def _probe(): Simulation as CoreSimulation, SimulationBuilder, ) + + # Capture the public API before country import installs compatibility. + from policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, + ) from policyengine_us import Microsimulation, Simulation from policyengine_us.system import system @@ -367,6 +373,136 @@ def set_input(self, variable, input_period, value): del simulation gc.collect() + # Subsampling rebuilds holder storage from a dataframe. Its new stores + # must retain the distinction between automatic and explicit sources. + sample_columns = { + **columns, + "household_weight__2024": [1.0], + "taxable_interest_income__2024": [40_000], + "state_code__2024": ["TX"], + } + simulation = CoreSimulation( + dataset=pd.DataFrame(sample_columns), + tax_benefit_system=system, + default_input_period=2024, + default_calculation_period=2024, + ) + simulation.set_input("social_security_retirement", 2024, [12_000]) + simulation.subsample(n=1, seed=0, time_period=2024) + simulation.set_input("social_security_retirement", 2024, [24_000]) + result["subsample_public_override"] = _snapshot( + simulation, (2024, 2026), INGRESS_OUTCOMES + ) + del simulation + gc.collect() + + # Main writes the legacy dump first. The current checkout restores those + # exact public-only files, rather than fabricating a legacy dump with its + # own compatibility hooks. A normal later setter supplies both names in + # the current checkout's own dump and needs automatic-source provenance. + dump_directory = Path(os.environ["PE_US_DIFFERENTIAL_DUMPS"]) + legacy_dump = dump_directory / "legacy_public" + inputs = extraction_situation() + del inputs["people"]["retiree"]["social_security_retirement"] + simulation = Simulation(situation=inputs, tax_benefit_system=system) + for component, amount in zip(COMPONENTS, AMOUNTS[0]): + simulation.set_input(component, 2024, [amount]) + if has_levers: + for component in COMPONENTS: + for name in (component, f"{component}_reported"): + assert ( + simulation._get_recorded_input_array(name, period(2024)) is not None + ) + own_dump = dump_directory / "both_names" + dump_simulation(simulation, own_dump) + else: + own_dump = legacy_dump + dump_simulation(simulation, legacy_dump) + + live = _snapshot(simulation, (2024, 2026), INGRESS_OUTCOMES) + restored = restore_simulation(legacy_dump, system) + result["dump_restore_legacy_public"] = _snapshot( + restored, (2024, 2026), INGRESS_OUTCOMES + ) + assert result["dump_restore_legacy_public"] == live + restored.set_input("social_security_retirement", 2024, [24_000]) + result["dump_restore_legacy_public_override"] = _snapshot( + restored, (2024, 2026), INGRESS_OUTCOMES + ) + del restored + gc.collect() + + restored = restore_simulation(own_dump, system) + result["dump_restore_both_names"] = _snapshot( + restored, (2024, 2026), INGRESS_OUTCOMES + ) + assert result["dump_restore_both_names"] == live + for candidate in (simulation, restored): + candidate.set_input("social_security_retirement", 2024, [24_000]) + result["dump_restore_both_names_override"] = _snapshot( + restored, (2024, 2026), INGRESS_OUTCOMES + ) + assert result["dump_restore_both_names_override"] == _snapshot( + simulation, (2024, 2026), INGRESS_OUTCOMES + ) + del simulation, restored + gc.collect() + + if has_levers: + # An explicit twin is intentionally independent of a later final + # public override. Restoring must not label every twin automatic. + simulation = Simulation( + situation=extraction_situation(), tax_benefit_system=system + ) + simulation.set_input("social_security_retirement_reported", 2024, [6_000]) + simulation.set_input("social_security_retirement", 2024, [12_000]) + explicit_dump = dump_directory / "explicit_reported" + dump_simulation(simulation, explicit_dump) + restored = restore_simulation(explicit_dump, system) + expected_future = np.asarray( + simulation.calculate("social_security_retirement", 2026) + ).copy() + for candidate in (simulation, restored): + candidate.set_input("social_security_retirement", 2024, [24_000]) + np.testing.assert_array_equal( + candidate.calculate("social_security_retirement", 2024), [24_000] + ) + np.testing.assert_array_equal( + candidate.calculate("social_security_retirement", 2026), + expected_future, + ) + assert _snapshot(restored, (2024, 2026), INGRESS_OUTCOMES) == _snapshot( + simulation, (2024, 2026), INGRESS_OUTCOMES + ) + del simulation, restored + gc.collect() + + # Core treats a nonderived raw cache value as an uprating source, + # without registering it for input replay. Dump restoration promotes + # it to a recorded input and must retain its automatic twin's origin. + inputs = extraction_situation() + del inputs["people"]["retiree"]["social_security_retirement"] + simulation = Simulation(situation=inputs, tax_benefit_system=system) + for component, amount in zip(COMPONENTS, AMOUNTS[0]): + simulation.get_holder(component).put_in_cache( + np.array([amount]), period(2024) + ) + for name in (component, f"{component}_reported"): + assert simulation._get_recorded_input_array(name, period(2024)) is None + raw_dump = dump_directory / "raw_cache" + dump_simulation(simulation, raw_dump) + restored = restore_simulation(raw_dump, system) + assert _snapshot(restored, (2024, 2026), INGRESS_OUTCOMES) == _snapshot( + simulation, (2024, 2026), INGRESS_OUTCOMES + ) + for candidate in (simulation, restored): + candidate.set_input("social_security_retirement", 2024, [24_000]) + assert _snapshot(restored, (2024, 2026), INGRESS_OUTCOMES) == _snapshot( + simulation, (2024, 2026), INGRESS_OUTCOMES + ) + del simulation, restored + gc.collect() + # Preserve Core's exact compound input and its annual/downstream # behavior. This simulation has no US wrapper. inputs = extraction_situation() @@ -553,11 +689,15 @@ def differential(tmp_path_factory): directory = tmp_path_factory.mktemp("social_security_differential") current = Path(__file__).resolve().parents[3] reference = _main_checkout(current, directory) - results = [] - for label, checkout in (("current", current), ("main", reference)): + results = {} + dump_directory = directory / "dumps" + dump_directory.mkdir() + # The main probe publishes a real legacy dump for the current probe. + for label, checkout in (("main", reference), ("current", current)): destination = directory / f"{label}.json" environment = dict(os.environ) environment["PYTHONPATH"] = str(checkout) + environment["PE_US_DIFFERENTIAL_DUMPS"] = str(dump_directory) for variable in ( "OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", @@ -585,8 +725,8 @@ def differential(tmp_path_factory): f"{label} differential probe failed:\n" f"{completed.stdout[-4000:]}\n{completed.stderr[-8000:]}" ) - results.append(json.loads(destination.read_text())) - return results + results[label] = json.loads(destination.read_text()) + return results["current"], results["main"] @pytest.mark.parametrize( @@ -620,6 +760,11 @@ def differential(tmp_path_factory): "direct_public_holder", "subclass_public_set_input_holder", "raw_core_dataset", + "subsample_public_override", + "dump_restore_legacy_public", + "dump_restore_legacy_public_override", + "dump_restore_both_names", + "dump_restore_both_names_override", "core_builder_exact_multiyear", "clone", "branch", diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index fc0df93e547..2b2b3de5ef1 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -6,6 +6,7 @@ """ import json +from types import SimpleNamespace import numpy as np import pandas as pd @@ -392,6 +393,191 @@ def test_holder_write_before_branch_creation_stays_in_its_named_branch(): assert amount(simulation, public, 2026) == 0 +@pytest.mark.parametrize("is_input", (None, True)) +def test_context_stores_synchronize_sources_in_the_effective_branch(is_input): + retirement, disability = SOCIAL_SECURITY_COMPONENTS[:2] + simulation = make_simulation(inputs={f"{disability}_reported": {2024: 0}}) + # Core restoration and custom input handlers use _set directly. Its + # default branch follows the active input context, including a branch + # that has no simulation object yet. + simulation._user_input_contexts.append("context_inputs") + try: + for public in (retirement, disability): + simulation.get_holder(public)._set( + period(2024), np.array([12_000]), is_input=is_input + ) + finally: + simulation._user_input_contexts.pop() + + branch = simulation.get_branch("context_inputs") + for public in (retirement, disability): + assert amount(branch, public, 2024) == 12_000 + assert amount(simulation, public, 2024) == 0 + assert (public, "context_inputs", period(2024)) in branch._user_input_keys + assert amount(branch, f"{retirement}_reported", 2024) == 12_000 + # An inherited explicit zero keeps precedence over a public input. + assert amount(branch, f"{disability}_reported", 2024) == 0 + for public in (retirement, disability): + branch.set_input(public, 2024, [24_000]) + drop_inherited_values(branch) + assert amount(branch, f"{retirement}_reported", 2024) == 24_000 + assert amount(branch, retirement, 2026) == pytest.approx(27_044.32, abs=0.01) + assert amount(branch, f"{disability}_reported", 2024) == 0 + assert amount(branch, disability, 2026) == 0 + assert amount(simulation, f"{retirement}_reported", 2024) == 0 + + +def test_explicit_noninput_store_keeps_its_branch_inside_an_input_context(): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + simulation = make_simulation() + simulation._user_input_contexts.append("context_inputs") + try: + simulation.get_holder(public)._set( + period(2024), np.array([12_000]), is_input=False + ) + finally: + simulation._user_input_contexts.pop() + for name in (public, reported): + assert set(simulation.get_holder(name).get_known_branch_periods()) == { + ("default", period(2024)) + } + assert input_periods(simulation, name) == set() + assert amount(simulation, name, 2024) == 12_000 + assert amount(simulation, public, 2026) == pytest.approx(13_522.16, abs=0.01) + + +def test_failing_period_handler_synchronizes_the_inputs_it_already_stored(monkeypatch): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + simulation = make_simulation() + branch = simulation.get_branch("partial_handler") + + def store_then_fail(holder, requested_period, array): + holder._set(requested_period.this_year, array) + raise RuntimeError("handler failed after storing its annual input") + + monkeypatch.setattr(system.variables[public], "set_input", store_then_fail) + with pytest.raises(RuntimeError, match="after storing"): + branch.get_holder(public).set_input( + period("2024-07"), np.array([12_000]), branch.branch_name + ) + assert branch._user_input_contexts == [] + assert amount(branch, public, 2024) == 12_000 + assert amount(branch, reported, 2024) == 12_000 + assert input_periods(branch, reported) == {"2024"} + assert amount(branch, public, 2026) == pytest.approx(13_522.16, abs=0.01) + branch.set_input(public, 2024, [24_000]) + drop_inherited_values(branch) + assert amount(branch, reported, 2024) == 24_000 + assert amount(branch, public, 2026) == pytest.approx(27_044.32, abs=0.01) + assert amount(simulation, public, 2024) == 0 + assert amount(simulation, reported, 2024) == 0 + + +@pytest.mark.parametrize("route", ("store_default", "store_noninput", "cache")) +def test_raw_nonderived_stores_keep_carry_sources_without_becoming_inputs(route): + retirement, disability = SOCIAL_SECURITY_COMPONENTS[:2] + simulation = make_simulation(inputs={f"{disability}_reported": {2024: 0}}) + for public in (retirement, disability): + holder = simulation.get_holder(public) + if route == "cache": + holder.put_in_cache(np.array([12_000]), period(2024), derived=False) + else: + holder._set( + period(2024), + np.array([12_000]), + is_input=False if route == "store_noninput" else None, + ) + assert amount(simulation, public, 2024) == 12_000 + assert input_periods(simulation, public) == set() + assert public not in simulation.input_variables + assert amount(simulation, f"{retirement}_reported", 2024) == 12_000 + assert input_periods(simulation, f"{retirement}_reported") == set() + assert f"{retirement}_reported" not in simulation.input_variables + assert amount(simulation, retirement, 2026) == pytest.approx(13_522.16, abs=0.01) + assert amount(simulation, f"{disability}_reported", 2024) == 0 + assert amount(simulation, disability, 2026) == 0 + # A later replayable input replaces its automatic cache source and still + # respects a separately supplied reported source. + for public in (retirement, disability): + simulation.set_input(public, 2024, [24_000]) + drop_inherited_values(simulation) + assert amount(simulation, f"{retirement}_reported", 2024) == 24_000 + assert input_periods(simulation, f"{retirement}_reported") == {"2024"} + assert amount(simulation, retirement, 2026) == pytest.approx(27_044.32, abs=0.01) + assert amount(simulation, f"{disability}_reported", 2024) == 0 + assert amount(simulation, disability, 2026) == 0 + + +def test_raw_memory_mirror_keeps_a_surviving_disk_source_automatic(tmp_path): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + simulation = make_simulation() + reported_holder = simulation.get_holder(reported) + reported_holder._disk_storage = reported_holder.create_disk_storage( + tmp_path, preserve=True + ) + reported_holder._on_disk_storable = True + simulation.memory_config = SimpleNamespace( + max_memory_occupation_pc=0, + priority_variables=system.variables, + variables_to_drop=(), + ) + stored_period = period(2024) + simulation.set_input(public, 2024, [12_000]) + assert reported_holder._memory_storage.get(stored_period) is None + np.testing.assert_array_equal( + reported_holder._disk_storage.get(stored_period), [12_000] + ) + + # A raw memory cache shadows the disk input without replacing its + # replayable tier or erasing that tier's automatic-source provenance. + reported_holder._on_disk_storable = False + simulation.get_holder(public).put_in_cache(np.array([18_000]), stored_period) + np.testing.assert_array_equal( + reported_holder._memory_storage.get(stored_period), [18_000] + ) + np.testing.assert_array_equal( + simulation._get_recorded_input_array(reported, stored_period), [12_000] + ) + simulation._invalidate_all_caches() + assert reported_holder._memory_storage.get(stored_period) is None + np.testing.assert_array_equal( + reported_holder._disk_storage.get(stored_period), [12_000] + ) + reported_holder._on_disk_storable = True + simulation.set_input(public, 2024, [24_000]) + np.testing.assert_array_equal( + reported_holder._disk_storage.get(stored_period), [24_000] + ) + assert amount(simulation, reported, 2024) == 24_000 + assert amount(simulation, public, 2026) == pytest.approx(27_044.32, abs=0.01) + + +def test_derived_public_caches_never_become_reported_carry_sources(): + retirement, disability = SOCIAL_SECURITY_COMPONENTS[:2] + simulation = make_simulation() + # A calculation can cache through either API, even within an input + # handler. Core explicitly excludes derived stores from input contexts. + simulation._user_input_contexts.append(simulation.branch_name) + try: + simulation.get_holder(retirement)._set( + period(2024), np.array([12_000]), derived=True + ) + simulation.get_holder(disability).put_in_cache( + np.array([12_000]), period(2024), derived=True + ) + finally: + simulation._user_input_contexts.pop() + for public in (retirement, disability): + assert amount(simulation, public, 2024) == 12_000 + assert input_periods(simulation, public) == set() + assert simulation.get_array(f"{public}_reported", 2024) is None + assert input_periods(simulation, f"{public}_reported") == set() + assert amount(simulation, public, 2026) == 0 + + @pytest.mark.parametrize("wrapper", WRAPPERS) def test_branch_explicit_twin_overrides_inherited_automatic_source_privately(wrapper): public = SOCIAL_SECURITY_COMPONENTS[0] diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index e57fead21e4..948c3f1dca0 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -8,6 +8,7 @@ import json from contextvars import ContextVar from functools import wraps +from pathlib import Path from policyengine_core.holders.helpers import get_stored_input from policyengine_core.periods import period as to_period @@ -28,12 +29,22 @@ _initial_social_security_ingress = ContextVar( "initial_social_security_ingress", default=False ) +_reloaded_social_security_provenance = ContextVar( + "reloaded_social_security_provenance", default=None +) +_DUMP_PROVENANCE_FILE = "social_security_input_provenance.json" _PROVENANCE = "_social_security_input_provenance" _AUTOMATIC_INPUTS = "_social_security_automatic_reported_inputs" -def _recorded_input_location(simulation, name, period, *, branch_names=None): - """Identify the actual supplied tier, including one hidden by a cache.""" +def _recorded_input_location( + simulation, name, period, *, branch_names=None, include_unrecorded=False +): + """Locate a supplied input, or optionally a legacy nonderived cache seed. + + Sequence stamps distinguish automatic sources from explicit replacements + even when the caller writes exactly the same amount. + """ holder = simulation.get_holder(name) period = holder._storage_period(period) input_keys = getattr(simulation, "_user_input_keys", ()) @@ -41,13 +52,20 @@ def _recorded_input_location(simulation, name, period, *, branch_names=None): branch_names = simulation._get_visible_branch_names() for branch_name in branch_names: key = (name, branch_name, period) - if key not in input_keys: - continue for storage_name, storage, separator in ( ("memory", holder._memory_storage, ":"), ("disk", holder._disk_storage, "_"), ): - if holder._stores_user_input(period, branch_name, storage_name): + if storage is None: + continue + recorded = key in input_keys and holder._stores_user_input( + period, branch_name, storage_name + ) + if recorded or ( + include_unrecorded + and storage.has(period, branch_name) + and not storage.is_derived(period, branch_name) + ): sequence = storage._sequence_numbers.get( f"{branch_name}{separator}{period}" ) @@ -57,12 +75,11 @@ def _recorded_input_location(simulation, name, period, *, branch_names=None): def _is_automatic_reported_input(simulation, name, period, *, branch_names=None): location = _recorded_input_location( - simulation, name, period, branch_names=branch_names - ) - return ( - location is not None - and getattr(simulation, _AUTOMATIC_INPUTS, {}).get(location[0]) == location[1] + simulation, name, period, branch_names=branch_names, include_unrecorded=True ) + return location is not None and location[1] in getattr( + simulation, _AUTOMATIC_INPUTS, {} + ).get(location[0], ()) def _mark_automatic_reported_input(simulation, name, period, *, branch_name=None): @@ -71,12 +88,19 @@ def _mark_automatic_reported_input(simulation, name, period, *, branch_name=None name, period, branch_names=None if branch_name is None else (branch_name,), + include_unrecorded=True, ) if location is not None: # Core clones simulation attributes shallowly. Never mutate a shared # provenance dictionary: each write belongs to this simulation alone. automatic = dict(getattr(simulation, _AUTOMATIC_INPUTS, {})) - automatic[location[0]] = location[1] + # A raw memory mirror can temporarily hide a surviving disk input. + # Keep each tier's origin so dropping that mirror reveals the original + # automatic disk source, rather than turning it into an explicit twin. + stamps = automatic.get(location[0], ()) + automatic[location[0]] = frozenset( + stamp for stamp in stamps if stamp[0] != location[1][0] + ) | {location[1]} setattr(simulation, _AUTOMATIC_INPUTS, automatic) @@ -84,7 +108,12 @@ def _automatic_reported_metadata(simulation, *, exported_periods=None): automatic = {} for public_name in SOCIAL_SECURITY_COMPONENTS: name = f"{public_name}_reported" - for period in simulation._get_set_input_periods(name): + # Dumping iterates this dictionary. Inspect existing holders without + # creating missing twins while Core is iterating over their values. + holder = simulation.persons._holders.get(name) + if holder is None: + continue + for period in holder.get_input_periods(simulation.branch_name): if ( exported_periods is not None and (name, str(period)) not in exported_periods @@ -118,6 +147,8 @@ def install_social_security_input_compatibility(): and dataset adapters migrate initial sources. Holder-level interception preserves subsequent public overrides from every supported input route: simulation setters, direct holder writes, and custom period helpers. + Nonderived cache seeds retain Core's carryover semantics without becoming + registered inputs; calculated (derived) caches never synchronize. """ from policyengine_core.simulations import Simulation as CoreSimulation from policyengine_core.simulations.simulation_builder import SimulationBuilder @@ -136,78 +167,51 @@ def tracks_public_input(holder): and not _initial_social_security_ingress.get() ) - def input_locations(holder, branch_name): - return { - period: location - for stored_branch, period in holder.get_known_branch_periods() - if stored_branch == branch_name - and ( - location := _recorded_input_location( - holder.simulation, - holder.variable.name, - period, - branch_names=(branch_name,), - ) - ) - is not None - } - - def synchronize_public_input(holder, period, branch_name): + def synchronize_public_input(holder, period, branch_name, *, is_input): simulation = holder.simulation variable_name = holder.variable.name period = holder._storage_period(period) - if ( - _recorded_input_location( - simulation, variable_name, period, branch_names=(branch_name,) - ) - is None - ): - return reported_name = f"{variable_name}_reported" reported_holder = simulation.get_holder(reported_name) readable_branches = reported_holder._readable_branches(branch_name) reported_location = _recorded_input_location( - simulation, reported_name, period, branch_names=readable_branches + simulation, + reported_name, + period, + branch_names=readable_branches, + include_unrecorded=True, ) if reported_location is None or _is_automatic_reported_input( simulation, reported_name, period, branch_names=readable_branches ): - supplied_value = get_stored_input(holder, period, branch_name) + supplied_value = ( + get_stored_input(holder, period, branch_name) + if is_input + else holder._get_array_from_storage(period, branch_name) + ) if supplied_value is not None: - reported_holder.set_input(period, supplied_value.copy(), branch_name) + if is_input: + reported_holder.set_input( + period, supplied_value.copy(), branch_name + ) + else: + # Core also uprates from legacy, nonderived cache stores. + # Mirror their source semantics without promoting them to + # replayable/exported inputs. Calculated caches never enter. + reported_holder._set( + period, supplied_value.copy(), branch_name, is_input=False + ) _mark_automatic_reported_input( simulation, reported_name, period, branch_name=branch_name ) + if not is_input: + return # Core clones share this legacy list. Replace it privately so a # later override is extracted without changing the source clone. simulation.input_variables = list( dict.fromkeys([*simulation.input_variables, variable_name, reported_name]) ) - original_set_input = Holder.set_input - if not getattr(original_set_input, "_social_security_compatibility", False): - - @wraps(original_set_input) - def set_input(holder, period, array, branch_name="default"): - if not tracks_public_input(holder): - return original_set_input(holder, period, array, branch_name) - - before = input_locations(holder, branch_name) - result = original_set_input(holder, period, array, branch_name) - # Period helpers can store a different period from the one passed - # in: a monthly input to a yearly variable stores a rolling year. - # Inspect actual inputs under the written branch, including a - # branch slot this simulation will read only after cloning. - for supplied_period, location in input_locations( - holder, branch_name - ).items(): - if location != before.get(supplied_period): - synchronize_public_input(holder, supplied_period, branch_name) - return result - - set_input._social_security_compatibility = True - Holder.set_input = set_input - original_store = Holder._set if not getattr(original_store, "_social_security_compatibility", False): @@ -222,15 +226,19 @@ def store_input( is_input=None, sequence_number=None, ): - # Core's direct period helpers and documented custom input stores - # can call _set(is_input=True) without entering Holder.set_input. - # Normal setter calls synchronize once, after their handler ends. - track_direct_input = ( - is_input - and not getattr( - getattr(holder, "simulation", None), "_user_input_contexts", None - ) - and tracks_public_input(holder) + track_source = not derived and tracks_public_input(holder) + # Mirror Core's classification and effective branch, including + # restore contexts and custom handlers that store before raising. + contexts = ( + getattr(holder.simulation, "_user_input_contexts", None) + if track_source + else None + ) + effective_input = bool(contexts) if is_input is None else is_input + stored_branch = ( + contexts[-1] + if effective_input and contexts and branch_name == "default" + else branch_name ) result = original_store( holder, @@ -242,8 +250,10 @@ def store_input( is_input=is_input, sequence_number=sequence_number, ) - if track_direct_input: - synchronize_public_input(holder, period, branch_name) + if track_source: + synchronize_public_input( + holder, period, stored_branch, is_input=effective_input + ) return result store_input._social_security_compatibility = True @@ -269,11 +279,107 @@ def build_from_dataset(simulation): # constructor or an entity/variable builder. Migrate before Core # creates its reform baseline, exactly as for situation ingress. migrate_social_security_inputs(simulation) + reloaded = _reloaded_social_security_provenance.get() + if reloaded is not None and reloaded[0] is simulation: + _restore_automatic_reported_metadata(simulation, reloaded[1]) return result build_from_dataset._social_security_compatibility = True CoreSimulation.build_from_dataset = build_from_dataset + original_subsample = CoreSimulation.subsample + if not getattr(original_subsample, "_social_security_compatibility", False): + + @wraps(original_subsample) + def subsample(simulation, *args, **kwargs): + if not getattr( + simulation.tax_benefit_system, + "_social_security_input_compatibility", + False, + ): + return original_subsample(simulation, *args, **kwargs) + # Restore stamps during dataset rebuilding, before Core replays + # inputs and clones the rebuilt reform baseline. + token = _reloaded_social_security_provenance.set( + (simulation, _automatic_reported_metadata(simulation)) + ) + try: + return original_subsample(simulation, *args, **kwargs) + finally: + _reloaded_social_security_provenance.reset(token) + + subsample._social_security_compatibility = True + CoreSimulation.subsample = subsample + + from policyengine_core.tools import simulation_dumper + + original_dump_holder = simulation_dumper._dump_holder + if not getattr(original_dump_holder, "_social_security_compatibility", False): + + @wraps(original_dump_holder) + def dump_holder(holder, directory): + result = original_dump_holder(holder, directory) + simulation = holder.simulation + name = holder.variable.name + if name in _REPORTED_COMPONENTS and getattr( + simulation.tax_benefit_system, + "_social_security_input_compatibility", + False, + ): + exported_periods = { + (name, str(period)) for period in holder.get_known_periods() + } + metadata = _automatic_reported_metadata( + simulation, exported_periods=exported_periods + ) + if metadata is not None: + path = Path(directory) / name / _DUMP_PROVENANCE_FILE + path.write_text(json.dumps(metadata)) + return result + + dump_holder._social_security_compatibility = True + simulation_dumper._dump_holder = dump_holder + + original_restore_holder = simulation_dumper._restore_holder + if not getattr(original_restore_holder, "_social_security_compatibility", False): + + @wraps(original_restore_holder) + def restore_holder( + simulation, variable, directory, inputs=False, calculated_number=None + ): + if not getattr( + simulation.tax_benefit_system, + "_social_security_input_compatibility", + False, + ): + return original_restore_holder( + simulation, variable, directory, inputs, calculated_number + ) + if not inputs and not getattr( + simulation, "_social_security_dump_inputs_restored", False + ): + # Core finishes the input pass before restoring any calculated + # arrays. Load all twins first, so directory order cannot turn + # an explicit twin into an automatic source or vice versa. + migrate_social_security_inputs(simulation) + for name in _REPORTED_COMPONENTS: + path = Path(directory) / name / _DUMP_PROVENANCE_FILE + if path.exists(): + _restore_automatic_reported_metadata( + simulation, json.loads(path.read_text()) + ) + simulation._social_security_dump_inputs_restored = True + token = _initial_social_security_ingress.set(True) + try: + return original_restore_holder( + simulation, variable, directory, inputs, calculated_number + ) + finally: + _initial_social_security_ingress.reset(token) + + restore_holder._social_security_compatibility = True + simulation_dumper._restore_holder = restore_holder + original_extract = CoreSimulation.extract_person if not getattr(original_extract, "_social_security_compatibility", False): @@ -343,11 +449,12 @@ def invalidate_all_caches(simulation): "_social_security_input_compatibility", False, ): - for key, stamp in getattr(simulation, _AUTOMATIC_INPUTS, {}).items(): + for key, stamps in getattr(simulation, _AUTOMATIC_INPUTS, {}).items(): name, branch_name, period = key - if _recorded_input_location( + location = _recorded_input_location( simulation, name, period, branch_names=(branch_name,) - ) == (key, stamp): + ) + if location is not None and location[1] in stamps: automatic.append(key) result = original_invalidate(simulation) # Core replay assigns new store sequence numbers. Replaying an From cd6d2988ca52738ecf7f788654fc30f8f38e43f2 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 02:16:32 -0400 Subject: [PATCH 32/76] Keep legacy dump fixtures compatible with removed empty holders --- .../test_social_security_current_law.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 136143dbfa3..932a7a64a70 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -148,6 +148,18 @@ def _probe(): from policyengine_us import Microsimulation, Simulation from policyengine_us.system import system + def dump_without_empty_holders(simulation, directory): + dump_simulation(simulation, directory) + # Core can create empty holders for variables removed between schemas. + # Preserve entities and every directory containing any array or metadata. + for holder_directory in Path(directory).iterdir(): + if ( + holder_directory.name != "__entities__" + and holder_directory.is_dir() + and not any(holder_directory.iterdir()) + ): + holder_directory.rmdir() + class HouseholdMicrosimulation(Microsimulation): default_dataset = None @@ -414,10 +426,10 @@ def set_input(self, variable, input_period, value): simulation._get_recorded_input_array(name, period(2024)) is not None ) own_dump = dump_directory / "both_names" - dump_simulation(simulation, own_dump) + dump_without_empty_holders(simulation, own_dump) else: own_dump = legacy_dump - dump_simulation(simulation, legacy_dump) + dump_without_empty_holders(simulation, legacy_dump) live = _snapshot(simulation, (2024, 2026), INGRESS_OUTCOMES) restored = restore_simulation(legacy_dump, system) @@ -457,7 +469,7 @@ def set_input(self, variable, input_period, value): simulation.set_input("social_security_retirement_reported", 2024, [6_000]) simulation.set_input("social_security_retirement", 2024, [12_000]) explicit_dump = dump_directory / "explicit_reported" - dump_simulation(simulation, explicit_dump) + dump_without_empty_holders(simulation, explicit_dump) restored = restore_simulation(explicit_dump, system) expected_future = np.asarray( simulation.calculate("social_security_retirement", 2026) @@ -490,7 +502,7 @@ def set_input(self, variable, input_period, value): for name in (component, f"{component}_reported"): assert simulation._get_recorded_input_array(name, period(2024)) is None raw_dump = dump_directory / "raw_cache" - dump_simulation(simulation, raw_dump) + dump_without_empty_holders(simulation, raw_dump) restored = restore_simulation(raw_dump, system) assert _snapshot(restored, (2024, 2026), INGRESS_OUTCOMES) == _snapshot( simulation, (2024, 2026), INGRESS_OUTCOMES From 6ff667b3c41b1f13429feff2206bece018c5f617 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 02:31:52 -0400 Subject: [PATCH 33/76] Synchronize raw Social Security seeds during initial loading --- changelog.d/ss-benefit-levers.added.md | 2 +- .../test_social_security_inputs.py | 57 +++++++++++++++++++ .../tools/social_security_inputs.py | 28 +++++---- 3 files changed, 75 insertions(+), 12 deletions(-) diff --git a/changelog.d/ss-benefit-levers.added.md b/changelog.d/ss-benefit-levers.added.md index 21001b7f1e4..2a55671bd0b 100644 --- a/changelog.d/ss-benefit-levers.added.md +++ b/changelog.d/ss-benefit-levers.added.md @@ -1 +1 @@ -Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs migrate on every ingress route; later public overrides refresh automatic sources, preserve explicit reported inputs, and survive extraction. Core dump/restore and subsampling preserve source provenance; direct input-context stores and legacy nonderived cache seeds retain future uprating. +Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs migrate on every ingress route; later public overrides refresh automatic sources, preserve explicit reported inputs, and survive extraction. Core dump/restore and subsampling preserve source provenance; direct input-context stores and legacy nonderived cache seeds retain future uprating, including stores made during initial loading. diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index 2b2b3de5ef1..3dab24e2cb4 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -510,6 +510,63 @@ def test_raw_nonderived_stores_keep_carry_sources_without_becoming_inputs(route) assert amount(simulation, disability, 2026) == 0 +@pytest.mark.parametrize("route", ("cache", "store_noninput")) +@pytest.mark.parametrize("ingress", ("builder", "situation", "dataset")) +def test_initial_raw_handler_stores_keep_sources_and_explicit_twin_priority( + monkeypatch, route, ingress +): + retirement, disability = SOCIAL_SECURITY_COMPONENTS[:2] + + def store_raw_source(holder, requested_period, array): + if route == "cache": + holder.put_in_cache(array, requested_period.this_year) + else: + holder._set(requested_period.this_year, array, is_input=False) + + for public in (retirement, disability): + monkeypatch.setattr(system.variables[public], "set_input", store_raw_source) + inputs = { + retirement: {"2024-07": 12_000}, + f"{disability}_reported": {2024: 0}, + disability: {"2024-07": 12_000}, + } + if ingress == "builder": + simulation = make_core_simulation(inputs, route="entities") + elif ingress == "situation": + simulation = CoreSimulation( + situation=situation(inputs), tax_benefit_system=system + ) + else: + entities = ("household", "tax_unit", "spm_unit", "family", "marital_unit") + frame = pd.DataFrame( + { + "person_id": [1], + **{f"person_{entity}_id": [1] for entity in entities}, + **{f"{entity}_id": [1] for entity in entities}, + **{ + f"{name}__{input_period}": [value] + for name, values in inputs.items() + for input_period, value in values.items() + }, + } + ) + simulation = CoreSimulation( + dataset=frame, + tax_benefit_system=system, + default_input_period=2024, + default_calculation_period=2024, + ) + for public in (retirement, disability): + assert amount(simulation, public, 2024) == 12_000 + assert input_periods(simulation, public) == set() + assert amount(simulation, f"{retirement}_reported", 2024) == 12_000 + assert input_periods(simulation, f"{retirement}_reported") == set() + assert amount(simulation, retirement, 2026) == pytest.approx(13_522.16, abs=0.01) + assert amount(simulation, f"{disability}_reported", 2024) == 0 + assert input_periods(simulation, f"{disability}_reported") == {"2024"} + assert amount(simulation, disability, 2026) == 0 + + def test_raw_memory_mirror_keeps_a_surviving_disk_source_automatic(tmp_path): public = SOCIAL_SECURITY_COMPONENTS[0] reported = f"{public}_reported" diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 948c3f1dca0..27cc8d02337 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -156,15 +156,10 @@ def install_social_security_input_compatibility(): def tracks_public_input(holder): simulation = getattr(holder, "simulation", None) - return ( - holder.variable.name in SOCIAL_SECURITY_COMPONENTS - and getattr( - getattr(simulation, "tax_benefit_system", None), - "_social_security_input_compatibility", - False, - ) - and hasattr(simulation, "input_variables") - and not _initial_social_security_ingress.get() + return holder.variable.name in SOCIAL_SECURITY_COMPONENTS and getattr( + getattr(simulation, "tax_benefit_system", None), + "_social_security_input_compatibility", + False, ) def synchronize_public_input(holder, period, branch_name, *, is_input): @@ -209,7 +204,13 @@ def synchronize_public_input(holder, period, branch_name, *, is_input): # Core clones share this legacy list. Replace it privately so a # later override is extracted without changing the source clone. simulation.input_variables = list( - dict.fromkeys([*simulation.input_variables, variable_name, reported_name]) + dict.fromkeys( + [ + *getattr(simulation, "input_variables", ()), + variable_name, + reported_name, + ] + ) ) original_store = Holder._set @@ -250,7 +251,12 @@ def store_input( is_input=is_input, sequence_number=sequence_number, ) - if track_source: + # Initial registered inputs wait until every explicit twin is + # loaded. Raw cache seeds never enter that migration, so mirror + # them immediately without making either value a supplied input. + if track_source and ( + not _initial_social_security_ingress.get() or not effective_input + ): synchronize_public_input( holder, period, stored_branch, is_input=effective_input ) From 7dce1d617dca037f87e79143bc8a2f0eba5e0f9d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 11:41:49 -0400 Subject: [PATCH 34/76] Refresh automatic benefit sources during composed loading --- changelog.d/ss-benefit-levers.added.md | 2 +- .../test_social_security_inputs.py | 27 +++++++- .../tools/social_security_inputs.py | 62 +++++++++++++++---- 3 files changed, 76 insertions(+), 15 deletions(-) diff --git a/changelog.d/ss-benefit-levers.added.md b/changelog.d/ss-benefit-levers.added.md index 2a55671bd0b..02add4ee9a0 100644 --- a/changelog.d/ss-benefit-levers.added.md +++ b/changelog.d/ss-benefit-levers.added.md @@ -1 +1 @@ -Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs migrate on every ingress route; later public overrides refresh automatic sources, preserve explicit reported inputs, and survive extraction. Core dump/restore and subsampling preserve source provenance; direct input-context stores and legacy nonderived cache seeds retain future uprating, including stores made during initial loading. +Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs migrate on every ingress route; later public overrides refresh automatic sources, preserve explicit reported inputs, and survive extraction. Core dump/restore and subsampling preserve source provenance; direct input-context stores and legacy nonderived cache seeds retain future uprating, including stores made during initial loading. Composed builder loading refreshes automatic sources for newly stored periods while preserving earlier public overrides and explicit twins. diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index 3dab24e2cb4..5c7fe98d02c 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -289,7 +289,32 @@ def test_repeated_public_writes_refresh_migrated_sources_and_future_years(wrappe override * 13_522.16 / 12_000 * 1.08, abs=0.01 ) assert input_periods(simulation, public) == {"2024"} - assert input_periods(simulation, reported) == {"2024"} + assert input_periods(simulation, reported) == {"2024"} + + +@pytest.mark.parametrize("explicit_source", (None, 0, 6_000)) +def test_existing_builder_refreshes_loaded_sources_and_keeps_earlier_overrides( + explicit_source, +): + inputs = {name: {2024: 12_000} for name in SOCIAL_SECURITY_COMPONENTS} + simulation = make_core_simulation(inputs) + for name in SOCIAL_SECURITY_COMPONENTS: + simulation.set_input(name, 2025, [18_000]) + if explicit_source is not None: + simulation.set_input(f"{name}_reported", 2024, [explicit_source]) + SimulationBuilder().build_from_variables( + system, + {name: {2024: [24_000]} for name in SOCIAL_SECURITY_COMPONENTS}, + simulation=simulation, + ) + for name in SOCIAL_SECURITY_COMPONENTS: + assert amount(simulation, f"{name}_reported", 2024) == ( + 24_000 if explicit_source is None else explicit_source + ) + assert amount(simulation, name, 2025) == 18_000 + assert input_periods(simulation, name) == ( + {"2025"} if explicit_source is None else {"2024", "2025"} + ) @pytest.mark.parametrize("wrapper", (*WRAPPERS, make_core_simulation)) diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 27cc8d02337..5a6ad5005cc 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -29,6 +29,9 @@ _initial_social_security_ingress = ContextVar( "initial_social_security_ingress", default=False ) +_initial_social_security_writes = ContextVar( + "initial_social_security_writes", default=None +) _reloaded_social_security_provenance = ContextVar( "reloaded_social_security_provenance", default=None ) @@ -251,6 +254,11 @@ def store_input( is_input=is_input, sequence_number=sequence_number, ) + loading_writes = _initial_social_security_writes.get() + if track_source and effective_input and loading_writes is not None: + loading_writes.setdefault(holder.simulation, {}).setdefault( + holder.variable.name, set() + ).add(holder._storage_period(period)) # Initial registered inputs wait until every explicit twin is # loaded. Raw cache seeds never enter that migration, so mirror # them immediately without making either value a supplied input. @@ -277,14 +285,21 @@ def build_from_dataset(simulation): ): return original_dataset_build(simulation) token = _initial_social_security_ingress.set(True) + writes = _initial_social_security_writes.get() + if writes is None: + writes = {} + writes_token = _initial_social_security_writes.set(writes) try: result = original_dataset_build(simulation) finally: + _initial_social_security_writes.reset(writes_token) _initial_social_security_ingress.reset(token) # Raw Core simulations use this loader without either country # constructor or an entity/variable builder. Migrate before Core # creates its reform baseline, exactly as for situation ingress. - migrate_social_security_inputs(simulation) + migrate_social_security_inputs( + simulation, loaded_public_inputs=writes.get(simulation, {}) + ) reloaded = _reloaded_social_security_provenance.get() if reloaded is not None and reloaded[0] is simulation: _restore_automatic_reported_metadata(simulation, reloaded[1]) @@ -376,11 +391,19 @@ def restore_holder( ) simulation._social_security_dump_inputs_restored = True token = _initial_social_security_ingress.set(True) + # Nested dict/entity/variable builders form one loading operation. + # Record actual stores, after period helpers have split them, so a + # builder on an existing simulation refreshes only its new inputs. + writes = _initial_social_security_writes.get() + if writes is None: + writes = {} + writes_token = _initial_social_security_writes.set(writes) try: return original_restore_holder( simulation, variable, directory, inputs, calculated_number ) finally: + _initial_social_security_writes.reset(writes_token) _initial_social_security_ingress.reset(token) restore_holder._social_security_compatibility = True @@ -518,7 +541,9 @@ def build(builder, tax_benefit_system, input_dict, simulation=None): ] ) ) - migrate_social_security_inputs(result) + migrate_social_security_inputs( + result, loaded_public_inputs=writes.get(result, {}) + ) _restore_automatic_reported_metadata(result, metadata) return result @@ -541,29 +566,40 @@ def build(builder, tax_benefit_system, input_dict, simulation=None): install_social_security_period_compatibility() -def migrate_social_security_inputs(simulation, *, include_baseline=True): +def migrate_social_security_inputs( + simulation, *, include_baseline=True, loaded_public_inputs=None +): """Move legacy component inputs, period by period, into pure sources. - A reported input in the same visible period makes the public input a - deliberate final override. Inputs in other periods do not affect that - choice. Recorded-input helpers exclude formula caches and honor visible - branch precedence. + An explicit twin in the same visible period keeps its priority. Automatic + twins refresh for public inputs stored by this loading operation. Earlier + public overrides are not initial inputs merely because another builder is + run on the same simulation. Without a loading record, migrate only legacy + public inputs with no twin. Recorded-input helpers exclude formula caches + and honor visible branch precedence. """ visible_branches = simulation._get_visible_branch_names() migrated = set() for public_name in SOCIAL_SECURITY_COMPONENTS: reported_name = f"{public_name}_reported" - reported_periods = set(simulation._get_set_input_periods(reported_name)) holder = simulation.get_holder(public_name) # Snapshot before deleting: deleting a multi-year period also deletes # contained supplied periods, including a shorter legacy source whose # lexical sort order happens to place it after the longer period. sources = {} for period in simulation._get_set_input_periods(public_name): - if period not in reported_periods: - value = simulation._get_recorded_input_array(public_name, period) - if value is not None: - sources[period] = value.copy() + location = _recorded_input_location( + simulation, reported_name, period, include_unrecorded=True + ) + if location is not None and ( + not _is_automatic_reported_input(simulation, reported_name, period) + or loaded_public_inputs is None + or period not in loaded_public_inputs.get(public_name, ()) + ): + continue + value = simulation._get_recorded_input_array(public_name, period) + if value is not None: + sources[period] = value.copy() if not sources: continue retained_overrides = [] @@ -571,7 +607,7 @@ def migrate_social_security_inputs(simulation, *, include_baseline=True): if ( name != public_name or branch_name not in visible_branches - or period not in reported_periods + or period in sources ): continue for storage_name, storage in ( From 9f8c6c304898489102cf05ac4d35894a569c9dfa Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 12:04:06 -0400 Subject: [PATCH 35/76] Reconcile successful benefit inputs when initial loading fails --- changelog.d/ss-benefit-levers.added.md | 2 +- .../test_social_security_current_law.py | 69 ++++++++++++++ .../test_social_security_inputs.py | 56 ++++++++++++ .../tools/social_security_inputs.py | 89 ++++++++++--------- 4 files changed, 174 insertions(+), 42 deletions(-) diff --git a/changelog.d/ss-benefit-levers.added.md b/changelog.d/ss-benefit-levers.added.md index 02add4ee9a0..f7ca7bf3844 100644 --- a/changelog.d/ss-benefit-levers.added.md +++ b/changelog.d/ss-benefit-levers.added.md @@ -1 +1 @@ -Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs migrate on every ingress route; later public overrides refresh automatic sources, preserve explicit reported inputs, and survive extraction. Core dump/restore and subsampling preserve source provenance; direct input-context stores and legacy nonderived cache seeds retain future uprating, including stores made during initial loading. Composed builder loading refreshes automatic sources for newly stored periods while preserving earlier public overrides and explicit twins. +Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs migrate on every ingress route; later public overrides refresh automatic sources, preserve explicit reported inputs, and survive extraction. Core dump/restore and subsampling preserve source provenance; direct input-context stores and legacy nonderived cache seeds retain future uprating, including stores made during initial loading. Composed builder loading refreshes automatic sources for newly stored periods while preserving earlier public overrides and explicit twins. Initial loading also reconciles completed stores when a later setter raises, preserving the original exception. diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 932a7a64a70..3682e1b5cb8 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -171,6 +171,14 @@ def set_input(self, variable, input_period, value): ) return super().set_input(variable, input_period, value) + class StoreThenRaiseSimulation(Simulation): + def set_input(self, variable, input_period, value): + result = super().set_input(variable, input_period, value) + error = getattr(self, "store_error", None) + if variable in COMPONENTS and error is not None: + raise error + return result + has_levers = "social_security_retirement_reported" in system.variables neutral = ( Reform.from_dict( @@ -385,6 +393,63 @@ def set_input(self, variable, input_period, value): del simulation gc.collect() + # An existing builder composes initial dataset ingress with a replacement + # public source. Copy before any future calculation, so cached uprating + # results cannot conceal a stale automatic reported source. + simulation = CoreSimulation( + dataset=pd.DataFrame(columns), + tax_benefit_system=system, + default_input_period=2024, + default_calculation_period=2024, + ) + SimulationBuilder().build_from_variables( + system, + {"social_security_retirement": {2024: [24_000]}}, + simulation=simulation, + ) + copied = simulation.clone(clone_tax_benefit_system=False) + copied.set_input("social_security_retirement", 2026, [18_000]) + SimulationBuilder().build_from_variables( + system, + {"social_security_retirement": {2024: [36_000]}}, + simulation=copied, + ) + result["dataset_existing_builder_overwrite"] = _snapshot( + simulation, (2024, 2026), INGRESS_OUTCOMES + ) + result["dataset_builder_clone_overwrite"] = _snapshot( + copied, (2024, 2026, 2027), INGRESS_OUTCOMES + ) + del simulation, copied + gc.collect() + + # Successful stores remain effective when a delegating setter raises. + # Exercise identical setters inside and outside builder loading and + # preserve the exact exception object supplied by the setter. + for route in ("builder_failed_initial_loading", "setter_failed_after_loading"): + inputs = extraction_situation() + del inputs["people"]["retiree"]["social_security_retirement"] + simulation = StoreThenRaiseSimulation( + situation=inputs, tax_benefit_system=system + ) + simulation.store_error = RuntimeError("setter failed after storing") + try: + if route == "builder_failed_initial_loading": + SimulationBuilder().build_from_variables( + system, + {"social_security_retirement": {2024: [24_000]}}, + simulation=simulation, + ) + else: + simulation.set_input("social_security_retirement", 2024, [24_000]) + except RuntimeError as error: + assert error is simulation.store_error + else: + raise AssertionError("The delegating setter did not raise") + result[route] = _snapshot(simulation, (2024, 2026), INGRESS_OUTCOMES) + del simulation + gc.collect() + # Subsampling rebuilds holder storage from a dataframe. Its new stores # must retain the distinction between automatic and explicit sources. sample_columns = { @@ -772,6 +837,10 @@ def differential(tmp_path_factory): "direct_public_holder", "subclass_public_set_input_holder", "raw_core_dataset", + "dataset_existing_builder_overwrite", + "dataset_builder_clone_overwrite", + "builder_failed_initial_loading", + "setter_failed_after_loading", "subsample_public_override", "dump_restore_legacy_public", "dump_restore_legacy_public_override", diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index 5c7fe98d02c..1202fb6a552 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -317,6 +317,62 @@ def test_existing_builder_refreshes_loaded_sources_and_keeps_earlier_overrides( ) +@pytest.mark.parametrize("route", ("builder", "dataset", "later")) +@pytest.mark.parametrize("explicit_source", (None, 0)) +def test_failed_loading_synchronizes_successful_stores_without_changing_error( + route, explicit_source +): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + failure = RuntimeError("setter failed after storing its input") + + class StoreThenRaise(CoreSimulation): + fail_after_store = False + + def set_input(self, name, input_period, value): + super().set_input(name, input_period, value) + if self.fail_after_store and name == public: + raise failure + + simulation = StoreThenRaise(situation=situation({}), tax_benefit_system=system) + if explicit_source is not None: + simulation.set_input(reported, 2024, [explicit_source]) + simulation.fail_after_store = True + with pytest.raises(RuntimeError) as caught: + if route == "builder": + SimulationBuilder().build_from_variables( + system, {public: {2024: [24_000]}}, simulation=simulation + ) + elif route == "dataset": + entities = ("household", "tax_unit", "spm_unit", "family", "marital_unit") + frame = pd.DataFrame( + { + "person_id": [1], + **{f"person_{entity}_id": [1] for entity in entities}, + **{f"{entity}_id": [1] for entity in entities}, + **( + {f"{reported}__2024": [explicit_source]} + if explicit_source is not None + else {} + ), + f"{public}__2024": [24_000], + } + ) + # Core accepts a DataFrame through its constructor, where it + # resolves the dataset object before invoking this same loader. + from policyengine_core.data import Dataset + + simulation.dataset = Dataset.from_dataframe(frame, time_period=2024) + simulation.build_from_dataset() + else: + simulation.set_input(public, 2024, [24_000]) + assert caught.value is failure + assert simulation._user_input_contexts == [] + assert amount(simulation, reported, 2024) == ( + 24_000 if explicit_source is None else explicit_source + ) + + @pytest.mark.parametrize("wrapper", (*WRAPPERS, make_core_simulation)) @pytest.mark.parametrize("explicit_source", (0, 12_000)) def test_direct_reported_holder_write_replaces_automatic_source_origin( diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 5a6ad5005cc..50c5aa54b18 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -284,6 +284,7 @@ def build_from_dataset(simulation): False, ): return original_dataset_build(simulation) + outermost = not _initial_social_security_ingress.get() token = _initial_social_security_ingress.set(True) writes = _initial_social_security_writes.get() if writes is None: @@ -294,15 +295,15 @@ def build_from_dataset(simulation): finally: _initial_social_security_writes.reset(writes_token) _initial_social_security_ingress.reset(token) - # Raw Core simulations use this loader without either country - # constructor or an entity/variable builder. Migrate before Core - # creates its reform baseline, exactly as for situation ingress. - migrate_social_security_inputs( - simulation, loaded_public_inputs=writes.get(simulation, {}) - ) - reloaded = _reloaded_social_security_provenance.get() - if reloaded is not None and reloaded[0] is simulation: - _restore_automatic_reported_metadata(simulation, reloaded[1]) + if outermost: + # A successful store remains an input even when a later + # setter raises. Reconcile before propagating that error. + migrate_social_security_inputs( + simulation, loaded_public_inputs=writes.get(simulation, {}) + ) + reloaded = _reloaded_social_security_provenance.get() + if reloaded is not None and reloaded[0] is simulation: + _restore_automatic_reported_metadata(simulation, reloaded[1]) return result build_from_dataset._social_security_compatibility = True @@ -391,19 +392,11 @@ def restore_holder( ) simulation._social_security_dump_inputs_restored = True token = _initial_social_security_ingress.set(True) - # Nested dict/entity/variable builders form one loading operation. - # Record actual stores, after period helpers have split them, so a - # builder on an existing simulation refreshes only its new inputs. - writes = _initial_social_security_writes.get() - if writes is None: - writes = {} - writes_token = _initial_social_security_writes.set(writes) try: return original_restore_holder( simulation, variable, directory, inputs, calculated_number ) finally: - _initial_social_security_writes.reset(writes_token) _initial_social_security_ingress.reset(token) restore_holder._social_security_compatibility = True @@ -513,38 +506,52 @@ def build(builder, tax_benefit_system, input_dict, simulation=None): for name, values in input_dict.items() if name != _PROVENANCE } + outermost = not _initial_social_security_ingress.get() token = _initial_social_security_ingress.set(True) + # Nested dict/entity/variable builders form one loading operation. + # Record actual stores after Core has handled their periods. + writes = _initial_social_security_writes.get() + if writes is None: + writes = {} + writes_token = _initial_social_security_writes.set(writes) + result = None try: result = original_build( builder, tax_benefit_system, input_dict, simulation ) finally: + _initial_social_security_writes.reset(writes_token) _initial_social_security_ingress.reset(token) - # Only build_from_dict rebuilds Core's legacy extraction list. - # Direct entity/variable builders need the same recorded inputs, - # including non-benefit sources and exact both-name overrides. - visible_branches = result._get_visible_branch_names() - recorded_names = { - name - for name, branch_name, _ in result._user_input_keys - if branch_name in visible_branches - } - result.input_variables = list( - dict.fromkeys( - [ - *getattr(result, "input_variables", []), - *sorted( + if outermost: + targets = dict(writes) + target = result if result is not None else simulation + if target is not None: + targets.setdefault(target, {}) + for target, loaded in targets.items(): + # Successful stores identify even a fresh build that + # raises before it can return its simulation. + visible_branches = target._get_visible_branch_names() + recorded_names = { name - for name in recorded_names - if result._get_set_input_periods(name) - ), - ] - ) - ) - migrate_social_security_inputs( - result, loaded_public_inputs=writes.get(result, {}) - ) - _restore_automatic_reported_metadata(result, metadata) + for name, branch_name, _ in target._user_input_keys + if branch_name in visible_branches + } + target.input_variables = list( + dict.fromkeys( + [ + *getattr(target, "input_variables", []), + *sorted( + name + for name in recorded_names + if target._get_set_input_periods(name) + ), + ] + ) + ) + migrate_social_security_inputs( + target, loaded_public_inputs=loaded + ) + _restore_automatic_reported_metadata(target, metadata) return result build._social_security_compatibility = True From 9f7db367f0cc10a8c4458d0652190a35860e0ef9 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 12:28:52 -0400 Subject: [PATCH 36/76] Search composed Social Security input routes with a state machine --- .../core/_social_security_state_backend.py | 771 ++++++++++++++++++ ...est_social_security_input_state_machine.py | 487 +++++++++++ 2 files changed, 1258 insertions(+) create mode 100644 policyengine_us/tests/core/_social_security_state_backend.py create mode 100644 policyengine_us/tests/core/test_social_security_input_state_machine.py diff --git a/policyengine_us/tests/core/_social_security_state_backend.py b/policyengine_us/tests/core/_social_security_state_backend.py new file mode 100644 index 00000000000..e59cadb616a --- /dev/null +++ b/policyengine_us/tests/core/_social_security_state_backend.py @@ -0,0 +1,771 @@ +"""Small storage backend for the Social Security input state machine. + +The country is intentionally never imported. Core's installed source bodies +implement every write, builder, copy, replay, extraction and dump route below; +the US module's complete function bodies implement the hooks. Only entities, +arrays, periods, storage and dependency bookkeeping are synthetic. This tests +input semantics, not benefit calculations or full-model performance. +""" + +from __future__ import annotations + +import ast +import copy +import itertools +import json +import logging +import os +import sys +import types +import warnings +from contextlib import ExitStack +from contextvars import ContextVar +from dataclasses import dataclass +from enum import Enum +from functools import lru_cache, wraps +from importlib.metadata import distribution +from pathlib import Path +from unittest.mock import patch + +import numpy as np + + +ROOT = Path(__file__).resolve().parents[2] +COMPONENTS = ( + "social_security_retirement", + "social_security_disability", + "social_security_survivors", + "social_security_dependents", +) + + +@dataclass(frozen=True, order=True) +class Instant: + year: int + month: int = 1 + + @property + def date(self): + return (self.year, self.month) + + +@dataclass(frozen=True) +class Period: + start: Instant + size: int = 1 + unit: str = "year" + + @property + def stop(self): + return Instant(self.start.year + self.size - 1, 12) + + def __str__(self): + if self.unit == "eternity": + return "eternity" + return ( + str(self.start.year) + if self.size == 1 + else f"year:{self.start.year}:{self.size}" + ) + + def contains(self, other): + return self.start <= other.start and self.stop >= other.stop + + +def as_period(value): + if isinstance(value, Period): + return value + if str(value) == "eternity": + return Period(Instant(0), unit="eternity") + parts = str(value).split(":") + return ( + Period(Instant(int(parts[-2])), int(parts[-1])) + if len(parts) == 3 + else Period(Instant(int(value))) + ) + + +@lru_cache(maxsize=None) +def source_tree(source): + return ast.parse(source.read_text()) + + +@lru_cache(maxsize=None) +def compiled_nodes(source, names, class_name=None): + """Cache code, never live functions or their mutable fixture namespaces.""" + tree = source_tree(source) + if class_name: + tree = next( + n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == class_name + ) + nodes = [ + next( + n + for n in tree.body + if isinstance(n, (ast.FunctionDef, ast.ClassDef)) and n.name == name + ) + for name in names + ] + future = ast.ImportFrom( + module="__future__", names=[ast.alias(name="annotations")], level=0 + ) + code = ast.fix_missing_locations(ast.Module(body=[future, *nodes], type_ignores=[])) + return compile(code, str(source), "exec") + + +def execute(source, names, namespace, class_name=None): + """Execute selected, unmodified source bodies in a fresh namespace.""" + exec(compiled_nodes(source, tuple(names), class_name), namespace) + return {name: namespace[name] for name in names} + + +@lru_cache(maxsize=None) +def compiled_hooks(source): + tree = source_tree(source) + body = [ + n + for n in tree.body + if isinstance(n, (ast.Assign, ast.AnnAssign, ast.FunctionDef)) + ] + code = ast.fix_missing_locations( + ast.Module( + body=[ + ast.ImportFrom( + module="__future__", names=[ast.alias(name="annotations")], level=0 + ), + *body, + ], + type_ignores=[], + ) + ) + return compile(code, str(source), "exec") + + +class Storage: + """Memory arrays plus Core's separate carry-over and input-tier marks.""" + + is_eternal = False + + def __init__(self): + self._arrays = {} + self._inputs = set() + self._sequence_numbers = {} + + def get(self, period, branch_name="default"): + return self._arrays.get(f"{branch_name}:{as_period(period)}") + + def has(self, period, branch_name="default"): + return f"{branch_name}:{as_period(period)}" in self._arrays + + def is_derived(self, period, branch_name="default"): + return ( + self.has(period, branch_name) + and f"{branch_name}:{as_period(period)}" not in self._inputs + ) + + def put( + self, value, period, branch_name="default", derived=False, sequence_number=None + ): + key = f"{branch_name}:{as_period(period)}" + self._arrays[key] = np.asarray(value).copy() + self._sequence_numbers[key] = sequence_number + if derived: + self._inputs.discard(key) + else: + self._inputs.add(key) + + def get_known_branch_periods(self): + return [ + (key.split(":", 1)[0], as_period(key.split(":", 1)[1])) + for key in self._arrays + ] + + def delete(self, period=None, branch_name="default"): + for key in list(self._arrays): + branch, when = key.split(":", 1) + if branch == branch_name and ( + period is None or as_period(period).contains(as_period(when)) + ): + del self._arrays[key] + self._unmark_dropped_keys() + + def mark_derived_except(self, kept): + self._inputs.intersection_update(kept) + + def drop_computed(self, since=None): + dropped = [ + key + for key in self._arrays + if key not in self._inputs + and (since is None or (self._sequence_numbers.get(key) or 0) >= since) + ] + for key in dropped: + del self._arrays[key] + self._unmark_dropped_keys() + return len(dropped) + + def _unmark_dropped_keys(self): + self._inputs.intersection_update(self._arrays) + + def _stop_sharing_dropped_keys(self): + pass + + def _forget_dropped_numbers(self): + self._sequence_numbers = { + key: number + for key, number in self._sequence_numbers.items() + if key in self._arrays + } + + def clone(self, **kwargs): + return copy.deepcopy(self) + + +class Disk: + """Tiny on-disk adapter used by unchanged public Core dump functions.""" + + def __init__(self, storage_dir, **kwargs): + self.storage_dir = str(storage_dir) + self.values = {} + Path(storage_dir).mkdir(exist_ok=True) + + def put(self, value, period): + Path(self.storage_dir, f"default_{period}.npy").write_text( + json.dumps(np.asarray(value).tolist()) + ) + + def restore(self): + self.values = { + as_period(path.stem.removeprefix("default_")): np.array( + json.loads(path.read_text()) + ) + for path in Path(self.storage_dir).glob("*.npy") + } + + def get_known_periods(self): + return list(self.values) + + def get(self, period): + return self.values[as_period(period)] + + +class History: + def copy(self): + return History() + + def record_unknown_sources(self, *args): + pass + + +class Backend: + """Install the unchanged hooks around unchanged Core routes in isolation.""" + + def __init__(self): + self.stack = ExitStack() + try: + self._setup() + except BaseException: + self.stack.close() + raise + + def _setup(self): + core = Path(distribution("policyengine-core").locate_file("policyengine_core")) + counter = itertools.count(1) + + def next_sequence(): + return next(counter) + + periods = types.SimpleNamespace( + period=as_period, + YEAR="year", + MONTH="month", + ETERNITY="eternity", + key_period_size=lambda p: p.size, + ) + commons = types.SimpleNamespace( + empty_clone=lambda value: object.__new__(type(value)) + ) + calculation_frames = ContextVar("state_machine_frames", default=()) + branch_clone = ContextVar("state_machine_branch_clone", default=None) + + class BranchClone: + def __init__(self, simulation): + self.simulation = simulation + self.pending = True + + class Encoder(json.JSONEncoder): + def default(self, value): + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + return super().default(value) + + @dataclass + class PreservedInput: + variable_name: str + branch_name: str + period: Period + value: object + storage: str + + person = types.SimpleNamespace(key="person", plural="people", is_person=True) + household = types.SimpleNamespace( + key="household", plural="households", is_person=False + ) + + class Holder: + def __init__(self, population, name): + self.population = population + self.simulation = population.simulation + self.variable = self.simulation.tax_benefit_system.get_variable(name) + self._memory_storage = Storage() + self._disk_storage = None + self._on_disk_storable = False + self._do_not_store = False + + def _storage_period(self, period): + return as_period(period) + + def _raise_if_input_contains_nan(self, array): + assert np.isfinite(array).all() + + def _to_array(self, value, validate_nan=False): + return np.asarray(value, dtype=float).reshape(1) + + def _record_store(self, *args): + pass + + def _record_inputs(self, *args): + pass + + def _evict_fast_cache(self, *args, **kwargs): + pass + + def _input_unchanged(self, *args): + return False + + def _stores(self, period, branch_name): + return self._memory_storage.has(period, branch_name) + + def _get_array_from_storage(self, period, branch_name="default"): + return self._memory_storage.get(period, branch_name) + + def get_array(self, period, branch_name="default"): + branch = self._branch_storing(period, branch_name) + return ( + None + if branch is None + else self._get_array_from_storage(period, branch) + ) + + def get_known_periods(self): + return list( + dict.fromkeys( + p for _, p in self._memory_storage.get_known_branch_periods() + ) + ) + + def create_disk_storage(self, directory, preserve=True): + return Disk(Path(directory) / self.variable.name) + + class Population: + def __init__(self, entity, simulation=None): + self.entity = entity + self.simulation = simulation + self._holders = {} + self.count = 1 + self.ids = np.array([1]) + self.members_entity_id = np.array([0]) + + def get_holder(self, name): + if not self.entity.is_person: + raise ValueError("All fixture variables belong to people") + if name not in self._holders: + self._holders[name] = Holder(self, name) + return self._holders[name] + + def clone(self, simulation, *args, **kwargs): + clone = Population(self.entity, simulation) + clone._holders = { + name: holder.clone(clone, **kwargs) + for name, holder in self._holders.items() + } + return clone + + class System: + _social_security_input_compatibility = True + person_entity = person + group_entities = [household] + cache_blacklist = [] + data_modified = False + + def __init__(self): + self.variables = { + name: types.SimpleNamespace( + name=name, + entity=person, + definition_period="year", + is_neutralized=False, + value_type=float, + set_input=None, + end=None, + ) + for name in ( + *COMPONENTS, + *(f"{name}_reported" for name in COMPONENTS), + ) + } + + def instantiate_entities(self): + return { + entity.key: Population(entity) for entity in (person, household) + } + + def get_variable(self, name, **kwargs): + return self.variables[name] + + def entities_plural(self): + return ["people", "households"] + + def clone(self): + return self + + system = System() + + class Simulation: + default_role = "member" + default_input_period = "2024" + default_calculation_period = "2024" + + def __init__(self, tax_benefit_system=None, populations=None, **kwargs): + self.tax_benefit_system = tax_benefit_system or system + self._user_input_keys = set() + self._user_input_contexts = [] + self._inputs_set = 0 + self._calculations_started = 0 + self.input_variables = [] + self.parent_branch = None + self.branch_name = "default" + self.branches = {} + self._data_storage_dir = None + self._storage_dir_parent = None + self.invalidated_caches = set() + self._fast_cache = {} + self.start_instant = None + self.opt_out_cache = False + self.trace = False + self.fail_after_store = False + self.build_from_populations( + populations or self.tax_benefit_system.instantiate_entities() + ) + + def build_from_populations(self, populations): + self.populations = populations + self.persons = populations["person"] + for population in populations.values(): + population.simulation = self + + def get_holder(self, name): + return self.persons.get_holder(name) + + def get_variable_population(self, name): + return self.persons + + def _drop_values_that_may_depend_on(self, *args): + pass + + def _get_store_history(self): + return History() + + def _get_requested_variables(self): + return set() + + def _drop_computed(self, since=None): + return sum( + holder._drop_computed(since) + for population in self.populations.values() + for holder in population._holders.values() + ) + + def subsample(self, *args): + raise NotImplementedError( + "Dataset subsampling is outside this storage backend" + ) + + def calculate(self, name, period): + return self.get_holder(name).get_array( + as_period(period), self.branch_name + ) + + def _calculate(self, name, period): + return self.calculate(name, period) + + class DelegatingSimulation(Simulation): + def set_input(self, name, period, value): + super().set_input(name, period, value) + if self.fail_after_store and name == self.fail_name: + self.last_error = RuntimeError( + "stored input before intentional setter failure" + ) + raise self.last_error + + class Builder: + default_period = None + + def __init__(self): + self.input_buffer = {} + self.entity_counts = {} + self.memberships = {} + + def explicit_singular_entities(self, system, values): + return values + + def register_variable(self, *args): + pass + + def add_person_entity(self, entity, values): + self.persons_plural = entity.plural + for inputs in values.values(): + for name, dated in inputs.items(): + self.input_buffer[name] = { + str(period): [amount] for period, amount in dated.items() + } + return list(values) + + def add_group_entity(self, *args): + pass + + def add_default_group_entity(self, *args): + pass + + def declare_person_entity(self, key, ids): + self.populations[key].ids = np.asarray(ids) + + def declare_entity(self, key, ids): + self.populations[key].ids = np.asarray(ids) + + def join_with_persons(self, *args): + pass + + class Dataset: + FLAT_FILE = "flat" + TIME_PERIOD_ARRAYS = "period_arrays" + data_format = TIME_PERIOD_ARRAYS + time_period = "2024" + + def __init__(self, values): + self.values = values + + def load(self): + return { + "person_id": {"2024": [1]}, + "household_id": {"2024": [1]}, + "person_household_id": {"2024": [1]}, + **self.values, + } + + namespace = dict( + __name__="policyengine_core.holders.holder", + __package__="policyengine_core.holders", + np=np, + numpy=np, + periods=periods, + Period=Period, + os=os, + warnings=warnings, + types=types, + json=json, + logging=logging, + commons=commons, + next_sequence_number=next_sequence, + _calculation_frames=calculation_frames, + _branch_clone=branch_clone, + _BranchClone=BranchClone, + _note_unobserved_activity=lambda: None, + directory_containing=lambda path: None, + PreservedUserInput=PreservedInput, + NpEncoder=Encoder, + ETERNITY="eternity", + Dataset=Dataset, + Simulation=Simulation, + SimulationBuilder=Builder, + _get_person_count=lambda values: 1, + check_type=lambda *args: None, + PeriodMismatchError=ValueError, + copy=copy, + ) + holder_methods = ( + "clone", + "set_input", + "_set", + "put_in_cache", + "_seed_input_storage", + "_record_input_storage", + "_stores_user_input", + "_readable_branches", + "_branch_storing", + "is_derived", + "get_input_periods", + "_drop_computed", + "delete_arrays", + "_stored_keys", + "_forget_deleted_inputs", + ) + for name, method in execute( + core / "holders/holder.py", holder_methods, namespace, "Holder" + ).items(): + setattr(Holder, name, method) + execute( + core / "holders/helpers.py", ("_is_input", "get_stored_input"), namespace + ) + sim_methods = ( + "set_input", + "build_from_dataset", + "clone", + "get_branch", + "apply_reform", + "_invalidate_all_caches", + "drop_computed_arrays", + "extract_person", + "_get_visible_branch_names", + "_get_set_input_periods", + "_get_recorded_input_array", + ) + for name, method in execute( + core / "simulations/simulation.py", sim_methods, namespace, "Simulation" + ).items(): + setattr(Simulation, name, method) + builder_methods = ( + "build_default_simulation", + "build_from_variables", + "build_from_dict", + "build_from_entities", + "finalize_variables_init", + ) + for name, method in execute( + core / "simulations/simulation_builder.py", + builder_methods, + namespace, + "SimulationBuilder", + ).items(): + setattr(Builder, name, method) + + def module(name, **values): + result = types.ModuleType(name) + result.__dict__.update(values) + return result + + dumper = module( + "policyengine_core.tools.simulation_dumper", + os=os, + warnings=warnings, + Simulation=Simulation, + OnDiskStorage=Disk, + ETERNITY="eternity", + next_sequence_number=next_sequence, + DERIVED_PERIODS_FILE="derived_periods.txt", + DUMP_FORMAT_FILE="dump_format.txt", + DUMP_FORMAT_VERSION=1, + _dump_entity=lambda *args: None, + _restore_entity=lambda *args: 1, + ) + execute( + core / "tools/simulation_dumper.py", + ( + "dump_simulation", + "restore_simulation", + "_dump_holder", + "_restore_holder", + "_restore_input", + ), + dumper.__dict__, + ) + helpers = module( + "policyengine_core.holders.helpers", + _is_input=namespace["_is_input"], + _period_from_storage_key=as_period, + get_stored_input=namespace["get_stored_input"], + ) + period_hooks = module( + "policyengine_us.tools.social_security_periods", + np=np, + YEAR="year", + CoreSimulation=Simulation, + SOCIAL_SECURITY_COMPONENTS=COMPONENTS, + wraps=wraps, + ) + enums = execute( + ROOT / "tools/social_security_benefit_levers.py", + ("SocialSecurityBenefitSource", "SocialSecurityFund"), + {"Enum": Enum, "__name__": __name__}, + ) + period_hooks.__dict__.update(enums) + execute( + ROOT / "tools/social_security_periods.py", + ( + "install_social_security_period_compatibility", + "_calculate_social_security_period", + ), + period_hooks.__dict__, + ) + modules = { + "policyengine_core": module("policyengine_core"), + "policyengine_core.holders": module( + "policyengine_core.holders", Holder=Holder + ), + "policyengine_core.holders.helpers": helpers, + "policyengine_core.simulations": module( + "policyengine_core.simulations", Simulation=Simulation + ), + "policyengine_core.simulations.simulation": module( + "policyengine_core.simulations.simulation", + _note_input_change=lambda *args: None, + ), + "policyengine_core.simulations.simulation_builder": module( + "policyengine_core.simulations.simulation_builder", + SimulationBuilder=Builder, + ), + "policyengine_core.tools": module( + "policyengine_core.tools", simulation_dumper=dumper + ), + "policyengine_core.tools.simulation_dumper": dumper, + "policyengine_us": module("policyengine_us", Simulation=Simulation), + "policyengine_us.tools": module("policyengine_us.tools"), + "policyengine_us.tools.social_security_periods": period_hooks, + } + self.stack.enter_context(patch.dict(sys.modules, modules)) + # Mutation evidence may supply a separate, minimally reverted source + # copy; normal tests always execute the actual production module. + hook_source = Path( + os.environ.get( + "SOCIAL_SECURITY_INPUT_SOURCE", ROOT / "tools/social_security_inputs.py" + ) + ) + hook_namespace = dict( + json=json, + ContextVar=ContextVar, + wraps=wraps, + Path=Path, + get_stored_input=namespace["get_stored_input"], + to_period=as_period, + NpEncoder=Encoder, + ETERNITY="eternity", + ) + execute( + ROOT / "tools/period_branch.py", ("delete_input_values",), hook_namespace + ) + # Keep every assignment and function body from the real hook module; + # execute anew for each example, so ContextVars and stamps stay private. + exec(compiled_hooks(hook_source), hook_namespace) + hook_namespace["install_social_security_input_compatibility"]() + self.Simulation = Simulation + self.DelegatingSimulation = DelegatingSimulation + self.Builder = Builder + self.Dataset = Dataset + self.system = system + self.hooks = hook_namespace + self.dumper = dumper + + def close(self): + self.stack.close() diff --git a/policyengine_us/tests/core/test_social_security_input_state_machine.py b/policyengine_us/tests/core/test_social_security_input_state_machine.py new file mode 100644 index 00000000000..5b8c7c37245 --- /dev/null +++ b/policyengine_us/tests/core/test_social_security_input_state_machine.py @@ -0,0 +1,487 @@ +"""Search compositions of public and reported Social Security input routes. + +Property: for every component and exact supplied period, an explicit reported +input has priority; otherwise the reported value follows the latest public +input, with automatic provenance. Copies preserve that property and public +final overrides, while subsequent writes remain isolated from their source. + +The backend executes installed Core source bodies and the production US hooks +without importing the country. Run this file directly for a country-free test: +``python policyengine_us/tests/core/test_social_security_input_state_machine.py``. +""" + +from __future__ import annotations + +import copy +import importlib.util +import sys +import tempfile +import unittest +from dataclasses import dataclass, field +from pathlib import Path + +from hypothesis import HealthCheck, settings, strategies as st +from hypothesis.stateful import RuleBasedStateMachine, invariant, rule + + +spec = importlib.util.spec_from_file_location( + "_social_security_state_backend", + Path(__file__).with_name("_social_security_state_backend.py"), +) +backend_module = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = backend_module +spec.loader.exec_module(backend_module) +Backend = backend_module.Backend +COMPONENTS = backend_module.COMPONENTS +as_period = backend_module.as_period + +COMPONENT = st.sampled_from(COMPONENTS) +PERIOD = st.sampled_from(("2023", "2024", "2025", "year:2023:2", "year:2024:2")) +AMOUNT = st.integers(min_value=0, max_value=48_000) +TARGET = st.integers(min_value=0, max_value=31) + + +@dataclass +class Input: + amount: int + registered: bool + route: str + + +@dataclass +class Model: + # These dictionaries describe caller intent, never inspect production + # provenance to decide whether a twin is automatic or explicit. + public: dict = field(default_factory=dict) + explicit: dict = field(default_factory=dict) + overrides: dict = field(default_factory=dict) + + def write( + self, component, reported, period, amount, registered, route, initial=False + ): + key = (component, period) + value = Input(amount, registered, route) + if reported: + self.explicit[key] = value + else: + self.public[key] = value + if initial and key not in self.explicit: + self.overrides.pop(key, None) + else: + self.overrides[key] = value + + def keep_registered(self): + self.public = { + key: value for key, value in self.public.items() if value.registered + } + self.explicit = { + key: value for key, value in self.explicit.items() if value.registered + } + self.overrides = { + key: value for key, value in self.overrides.items() if value.registered + } + + def dump_promotes_raw(self): + for mapping in (self.public, self.explicit, self.overrides): + for value in mapping.values(): + value.registered = True + + +class SocialSecurityInputsMachine(RuleBasedStateMachine): + """All writes may slot earlier clones, branches, or restored copies.""" + + def __init__(self): + super().__init__() + self.backend = Backend() + self.states = [(self.backend.DelegatingSimulation(), Model())] + self.temp = tempfile.TemporaryDirectory(prefix="ss-input-machine-") + self.dump_count = 0 + + def teardown(self): + self.temp.cleanup() + self.backend.close() + + def state(self, slot): + return self.states[slot % len(self.states)] + + @staticmethod + def name(component, reported): + return f"{component}_reported" if reported else component + + def append(self, simulation, model): + self.states.append((simulation, model)) + + @staticmethod + def write_key(model, component, period, data): + # Reuse caller-known slots as well as drawing fresh ones. This makes + # collisions between independent routes, copies and replay common. + keys = sorted(model.public.keys() | model.explicit.keys()) + if ( + data is not None + and keys + and data.draw(st.booleans(), label="reuse existing input") + ): + return data.draw( + st.sampled_from(keys), label="existing component and period" + ) + return component, period + + @rule( + slot=TARGET, + component=COMPONENT, + reported=st.booleans(), + period=PERIOD, + amount=AMOUNT, + data=st.data(), + route=st.sampled_from(("simulation", "holder", "delegating")), + ) + def set_input(self, slot, component, reported, period, amount, route, data=None): + simulation, model = self.state(slot) + component, period = self.write_key(model, component, period, data) + name = self.name(component, reported) + if route == "holder": + simulation.get_holder(name).set_input( + as_period(period), [amount], simulation.branch_name + ) + else: + # All live instances are delegating subclasses or genuine Core + # copies thereof; explicitly invoke Core for its separate route. + setter = ( + self.backend.Simulation.set_input + if route == "simulation" + else type(simulation).set_input + ) + setter(simulation, name, period, [amount]) + model.write(component, reported, period, amount, True, route) + + @rule( + slot=TARGET, + component=COMPONENT, + reported=st.booleans(), + period=PERIOD, + amount=AMOUNT, + data=st.data(), + is_input=st.sampled_from((None, False, True)), + in_context=st.booleans(), + default_branch_argument=st.booleans(), + ) + def low_level_store( + self, + slot, + component, + reported, + period, + amount, + is_input, + in_context, + default_branch_argument, + data=None, + ): + simulation, model = self.state(slot) + component, period = self.write_key(model, component, period, data) + if in_context: + simulation._user_input_contexts.append(simulation.branch_name) + try: + # Core redirects default to an active input context for registered + # stores. Exercise both implicit and explicit branch selection. + branch_argument = ( + "default" + if default_branch_argument and in_context and is_input is not False + else simulation.branch_name + ) + simulation.get_holder(self.name(component, reported))._set( + as_period(period), [amount], branch_argument, is_input=is_input + ) + finally: + if in_context: + simulation._user_input_contexts.pop() + registered = in_context if is_input is None else is_input + model.write( + component, + reported, + period, + amount, + registered, + f"_set({is_input}, context={in_context})", + ) + + @rule( + slot=TARGET, + component=COMPONENT, + reported=st.booleans(), + period=PERIOD, + amount=AMOUNT, + data=st.data(), + ) + def nonderived_cache(self, slot, component, reported, period, amount, data=None): + simulation, model = self.state(slot) + component, period = self.write_key(model, component, period, data) + simulation.get_holder(self.name(component, reported)).put_in_cache( + [amount], as_period(period), simulation.branch_name, derived=False + ) + model.write(component, reported, period, amount, False, "put_in_cache") + + @rule( + slot=TARGET, + component=COMPONENT, + reported=st.booleans(), + period=PERIOD, + amount=AMOUNT, + data=st.data(), + derived=st.booleans(), + ) + def cache_result( + self, slot, component, reported, period, amount, derived, data=None + ): + """Calculated values must not acquire reported-source provenance.""" + if not derived: + return + simulation, model = self.state(slot) + component, period = self.write_key(model, component, period, data) + simulation.get_holder(self.name(component, reported)).put_in_cache( + [amount], as_period(period), simulation.branch_name, derived=True + ) + + @rule( + slot=TARGET, + component=COMPONENT, + reported=st.booleans(), + period=PERIOD, + amount=AMOUNT, + data=st.data(), + ) + def dataset_load(self, slot, component, reported, period, amount, data=None): + simulation, model = self.state(slot) + component, period = self.write_key(model, component, period, data) + simulation.dataset = self.backend.Dataset( + {self.name(component, reported): {period: [amount]}} + ) + simulation.build_from_dataset() + model = Model() + model.write(component, reported, period, amount, True, "dataset", initial=True) + self.states[slot % len(self.states)] = (simulation, model) + + @rule( + slot=TARGET, + component=COMPONENT, + reported=st.booleans(), + period=PERIOD, + amount=AMOUNT, + data=st.data(), + fresh=st.booleans(), + yaml=st.booleans(), + ) + def builder_load( + self, slot, component, reported, period, amount, fresh, yaml, data=None + ): + simulation, model = self.state(slot) + component, period = self.write_key(model, component, period, data) + if fresh: + simulation, model = None, Model() + name = self.name(component, reported) + builder = self.backend.Builder() + if yaml: + # This executes Core build_from_dict -> build_from_entities -> + # finalize_variables_init -> Holder.set_input, including nested + # loading contexts and exact multi-year periods. + inputs = { + "people": {"person_1": {name: {period: amount}}}, + "households": {"household": {"members": ["person_1"]}}, + } + result = builder.build_from_dict( + self.backend.system, inputs, simulation=simulation + ) + else: + result = builder.build_from_variables( + self.backend.system, {name: {period: [amount]}}, simulation=simulation + ) + if simulation is not None: + assert result is simulation + model.write( + component, + reported, + period, + amount, + True, + "yaml" if yaml else "builder", + initial=True, + ) + if fresh: + self.append(result, model) + + @rule( + slot=TARGET, + component=COMPONENT, + reported=st.booleans(), + period=PERIOD, + amount=AMOUNT, + data=st.data(), + initial=st.booleans(), + ) + def setter_stores_then_raises( + self, slot, component, reported, period, amount, initial, data=None + ): + simulation, model = self.state(slot) + component, period = self.write_key(model, component, period, data) + # A Core-created/restored instance is upgraded only by calling the + # delegating method explicitly, without replacing any hook. + original_class = type(simulation) + simulation.__class__ = self.backend.DelegatingSimulation + simulation.fail_after_store = True + simulation.fail_name = self.name(component, reported) + error = None + try: + name = self.name(component, reported) + if initial: + self.backend.Builder().build_from_variables( + self.backend.system, + {name: {period: [amount]}}, + simulation=simulation, + ) + else: + simulation.set_input(name, period, [amount]) + except RuntimeError as caught: + error = caught + finally: + simulation.fail_after_store = False + simulation.__class__ = original_class + assert ( + error is not None + and str(error) == "stored input before intentional setter failure" + ) + assert error is simulation.last_error + model.write( + component, + reported, + period, + amount, + True, + "failed initial setter" if initial else "failed setter", + initial=initial, + ) + + @rule(slot=TARGET, branch=st.booleans()) + def copy_simulation(self, slot, branch): + simulation, model = self.state(slot) + copied = ( + simulation.get_branch(f"copy_{len(self.states)}") + if branch + else simulation.clone() + ) + self.append(copied, copy.deepcopy(model)) + + @rule(slot=TARGET) + def dump_restore(self, slot): + simulation, model = self.state(slot) + directory = Path(self.temp.name) / f"dump_{self.dump_count}" + self.dump_count += 1 + self.backend.dumper.dump_simulation(simulation, directory) + restored = self.backend.dumper.restore_simulation( + directory, self.backend.system + ) + copied_model = copy.deepcopy(model) + copied_model.dump_promotes_raw() + self.append(restored, copied_model) + + @rule(slot=TARGET) + def extract_reload(self, slot): + simulation, model = self.state(slot) + payload = simulation.extract_person() + restored = self.backend.Builder().build_from_dict(self.backend.system, payload) + copied_model = copy.deepcopy(model) + copied_model.keep_registered() + self.append(restored, copied_model) + + @rule(slot=TARGET, replay=st.booleans()) + def reform_or_drop(self, slot, replay): + simulation, model = self.state(slot) + if replay: + + class IdentityReform: + def apply(self, system): + pass + + simulation.apply_reform(IdentityReform()) + else: + simulation.drop_computed_arrays() + # Branches already made are replayed by Core apply_reform too, with + # their independent snapshots. A plain drop affects only its slot. + affected = {simulation} + if replay: + pending = list(simulation.branches.values()) + while pending: + branch = pending.pop() + affected.add(branch) + pending.extend(branch.branches.values()) + for other, other_model in self.states: + if other in affected: + other_model.keep_registered() + + @invariant() + def reported_values_and_provenance_follow_intent(self): + for simulation, model in self.states: + assert simulation._user_input_contexts == [] + for key in model.public.keys() | model.explicit.keys(): + component, period = key + expected = model.explicit.get(key, model.public.get(key)) + reported = f"{component}_reported" + actual = simulation.get_holder(reported).get_array( + as_period(period), simulation.branch_name + ) + assert actual is not None, (key, expected.route, "missing twin") + assert float(actual[0]) == expected.amount, ( + key, + expected.route, + float(actual[0]), + expected.amount, + ) + automatic = self.backend.hooks["_is_automatic_reported_input"]( + simulation, reported, as_period(period) + ) + assert automatic == (key not in model.explicit), ( + key, + expected.route, + "provenance", + automatic, + ) + for component in COMPONENTS: + reported_holder = simulation.get_holder(f"{component}_reported") + for period in reported_holder.get_input_periods(simulation.branch_name): + key = (component, str(period)) + assert key in model.public or key in model.explicit, ( + key, + "unexplained reported input", + ) + holder = simulation.get_holder(component) + for period in holder.get_input_periods(simulation.branch_name): + key = (component, str(period)) + expected = model.overrides.get(key) + assert expected is not None, (key, "unexpected public override") + assert ( + float(holder.get_array(period, simulation.branch_name)[0]) + == expected.amount + ) + for (component, period), expected in model.overrides.items(): + actual = simulation.get_holder(component).get_array( + as_period(period), simulation.branch_name + ) + assert actual is not None and float(actual[0]) == expected.amount, ( + (component, period), + expected.route, + "lost public override", + ) + + +TestSocialSecurityInputStateMachine = SocialSecurityInputsMachine.TestCase +TestSocialSecurityInputStateMachine.settings = settings( + max_examples=60, + stateful_step_count=25, + deadline=None, + derandomize=True, + report_multiple_bugs=False, + suppress_health_check=(HealthCheck.too_slow,), +) + + +if __name__ == "__main__": + unittest.main() From 8396422945b07b31448f708654f6aed14be9ac70 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 13:18:18 -0700 Subject: [PATCH 37/76] Resolve reported benefit sources at read time --- changelog.d/ss-benefit-levers.added.md | 2 +- .../core/_social_security_state_backend.py | 50 +- ...est_social_security_input_state_machine.py | 174 ++- .../test_social_security_current_law.py | 14 + .../test_social_security_inputs.py | 99 +- .../tools/social_security_inputs.py | 1364 ++++++++++------- 6 files changed, 1045 insertions(+), 658 deletions(-) diff --git a/changelog.d/ss-benefit-levers.added.md b/changelog.d/ss-benefit-levers.added.md index f7ca7bf3844..31b9fa30ce7 100644 --- a/changelog.d/ss-benefit-levers.added.md +++ b/changelog.d/ss-benefit-levers.added.md @@ -1 +1 @@ -Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs migrate on every ingress route; later public overrides refresh automatic sources, preserve explicit reported inputs, and survive extraction. Core dump/restore and subsampling preserve source provenance; direct input-context stores and legacy nonderived cache seeds retain future uprating, including stores made during initial loading. Composed builder loading refreshes automatic sources for newly stored periods while preserving earlier public overrides and explicit twins. Initial loading also reconciles completed stores when a later setter raises, preserving the original exception. +Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs remain independent sources, exposed through a read-time reported view without automatic twin arrays or provenance stamps. Initial loading routes sources so same-year benefit factors apply; later public writes remain final overrides. Explicit reported inputs retain priority, while exact compound totals keep their period semantics. Builders, partial loading failures, copies, cache replay, extraction, dump restoration and subsampling preserve successful source inputs. A state machine searches combinations of write and copy routes. diff --git a/policyengine_us/tests/core/_social_security_state_backend.py b/policyengine_us/tests/core/_social_security_state_backend.py index e59cadb616a..bd055237505 100644 --- a/policyengine_us/tests/core/_social_security_state_backend.py +++ b/policyengine_us/tests/core/_social_security_state_backend.py @@ -125,7 +125,7 @@ def compiled_hooks(source): body = [ n for n in tree.body - if isinstance(n, (ast.Assign, ast.AnnAssign, ast.FunctionDef)) + if isinstance(n, (ast.Assign, ast.AnnAssign, ast.ClassDef, ast.FunctionDef)) ] code = ast.fix_missing_locations( ast.Module( @@ -180,6 +180,11 @@ def get_known_branch_periods(self): for key in self._arrays ] + def get_known_periods(self): + return list( + dict.fromkeys(period for _, period in self.get_known_branch_periods()) + ) + def delete(self, period=None, branch_name="default"): for key in list(self._arrays): branch, when = key.split(":", 1) @@ -256,6 +261,9 @@ def copy(self): def record_unknown_sources(self, *args): pass + def record_derived(self, *args): + pass + class Backend: """Install the unchanged hooks around unchanged Core routes in isolation.""" @@ -410,6 +418,7 @@ def __init__(self): value_type=float, set_input=None, end=None, + is_input_variable=lambda name=name: name.endswith("_reported"), ) for name in ( *COMPONENTS, @@ -459,6 +468,12 @@ def __init__(self, tax_benefit_system=None, populations=None, **kwargs): self.build_from_populations( populations or self.tax_benefit_system.instantiate_entities() ) + if kwargs.get("situation") is not None: + # Country constructor plumbing uses the same actual Core + # builder; no input migration logic lives in this fixture. + Builder().build_from_dict( + self.tax_benefit_system, kwargs["situation"], simulation=self + ) def build_from_populations(self, populations): self.populations = populations @@ -577,6 +592,7 @@ def load(self): types=types, json=json, logging=logging, + log=logging.getLogger(__name__), commons=commons, next_sequence_number=next_sequence, _calculation_frames=calculation_frames, @@ -593,10 +609,15 @@ def load(self): _get_person_count=lambda values: 1, check_type=lambda *args: None, PeriodMismatchError=ValueError, + _period_from_storage_key=as_period, copy=copy, ) holder_methods = ( "clone", + "_get_array_from_storage", + "get_array", + "get_known_periods", + "get_known_branch_periods", "set_input", "_set", "put_in_cache", @@ -616,9 +637,16 @@ def load(self): core / "holders/holder.py", holder_methods, namespace, "Holder" ).items(): setattr(Holder, name, method) - execute( - core / "holders/helpers.py", ("_is_input", "get_stored_input"), namespace + helper_source = core / "holders/helpers.py" + # Later compatible Core releases add helper calls to these routes. + # Execute the installed helper bodies together so their internal + # dependencies remain the real implementation across Core versions. + helper_names = tuple( + node.name + for node in source_tree(helper_source).body + if isinstance(node, ast.FunctionDef) ) + core_helpers = execute(helper_source, helper_names, namespace) sim_methods = ( "set_input", "build_from_dataset", @@ -683,9 +711,8 @@ def module(name, **values): ) helpers = module( "policyengine_core.holders.helpers", - _is_input=namespace["_is_input"], _period_from_storage_key=as_period, - get_stored_input=namespace["get_stored_input"], + **core_helpers, ) period_hooks = module( "policyengine_us.tools.social_security_periods", @@ -742,21 +769,28 @@ def module(name, **values): "SOCIAL_SECURITY_INPUT_SOURCE", ROOT / "tools/social_security_inputs.py" ) ) - hook_namespace = dict( + hook_module = module( + "_social_security_state_hooks", json=json, ContextVar=ContextVar, + dataclass=dataclass, wraps=wraps, Path=Path, + next_sequence_number=next_sequence, get_stored_input=namespace["get_stored_input"], to_period=as_period, NpEncoder=Encoder, ETERNITY="eternity", ) + self.stack.enter_context( + patch.dict(sys.modules, {hook_module.__name__: hook_module}) + ) + hook_namespace = hook_module.__dict__ execute( ROOT / "tools/period_branch.py", ("delete_input_values",), hook_namespace ) - # Keep every assignment and function body from the real hook module; - # execute anew for each example, so ContextVars and stamps stay private. + # Keep every assignment, class and function from the real hook module; + # execute anew for each example so all mutable state stays private. exec(compiled_hooks(hook_source), hook_namespace) hook_namespace["install_social_security_input_compatibility"]() self.Simulation = Simulation diff --git a/policyengine_us/tests/core/test_social_security_input_state_machine.py b/policyengine_us/tests/core/test_social_security_input_state_machine.py index 5b8c7c37245..56a9e40bdf4 100644 --- a/policyengine_us/tests/core/test_social_security_input_state_machine.py +++ b/policyengine_us/tests/core/test_social_security_input_state_machine.py @@ -1,9 +1,9 @@ """Search compositions of public and reported Social Security input routes. -Property: for every component and exact supplied period, an explicit reported -input has priority; otherwise the reported value follows the latest public -input, with automatic provenance. Copies preserve that property and public -final overrides, while subsequent writes remain isolated from their source. +Property: a reported source uses an explicit input at or before its requested +period; otherwise its read-time view follows the latest public input. Exact +compound totals keep their own period semantics. Copies preserve those +sources and public final overrides, while later writes remain isolated. The backend executes installed Core source bodies and the production US hooks without importing the country. Run this file directly for a country-free test: @@ -50,45 +50,110 @@ class Input: @dataclass class Model: - # These dictionaries describe caller intent, never inspect production - # provenance to decide whether a twin is automatic or explicit. + # Caller inputs are distinct slots on each branch. A raw branch write + # hides, rather than erases, an inherited registered input. public: dict = field(default_factory=dict) explicit: dict = field(default_factory=dict) overrides: dict = field(default_factory=dict) + branch_name: str = "default" + visible_branches: tuple = ("default",) + + def keys(self): + return { + (component, period) + for mapping in (self.public, self.explicit, self.overrides) + for branch, component, period in mapping + if branch in self.visible_branches + } + + def visible(self, mapping, component, period): + return next( + ( + mapping[(branch, component, period)] + for branch in self.visible_branches + if (branch, component, period) in mapping + ), + None, + ) + + def reported(self, component, period): + explicit = self.visible(self.explicit, component, period) + if explicit is not None: + return explicit, False + requested = as_period(period) + for other_component, known in self.keys(): + candidate = as_period(known) + if ( + other_component == component + and ( + candidate.start <= requested.start + if requested.size == 1 + else candidate == requested + ) + and self.visible(self.explicit, component, known) is not None + ): + # Core uprates from the earlier explicit source. It must not + # see a same-period public fallback as a newer reported input. + return None, False + return self.visible(self.public, component, period), True def write( - self, component, reported, period, amount, registered, route, initial=False + self, + component, + reported, + period, + amount, + registered, + route, + initial=False, + branch_name=None, ): - key = (component, period) + branch_name = self.branch_name if branch_name is None else branch_name + key = (branch_name, component, period) value = Input(amount, registered, route) if reported: self.explicit[key] = value else: self.public[key] = value - if initial and key not in self.explicit: - self.overrides.pop(key, None) + if initial and self.visible(self.explicit, component, period) is None: + # Initial migration removes visible public slots, including + # a superseded inherited final override for the same period. + for branch in self.visible_branches: + self.overrides.pop((branch, component, period), None) else: self.overrides[key] = value def keep_registered(self): - self.public = { - key: value for key, value in self.public.items() if value.registered - } - self.explicit = { - key: value for key, value in self.explicit.items() if value.registered - } - self.overrides = { - key: value for key, value in self.overrides.items() if value.registered - } + for name in ("public", "explicit", "overrides"): + setattr( + self, + name, + { + key: value + for key, value in getattr(self, name).items() + if value.registered + }, + ) - def dump_promotes_raw(self): - for mapping in (self.public, self.explicit, self.overrides): - for value in mapping.values(): - value.registered = True + def flatten(self, promote_raw=False): + keys = self.keys() + for name in ("public", "explicit", "overrides"): + mapping = getattr(self, name) + flat = {} + for component, period in keys: + value = self.visible(mapping, component, period) + if value is not None: + value = copy.deepcopy(value) + if promote_raw: + value.registered = True + flat[("default", component, period)] = value + setattr(self, name, flat) + self.branch_name = "default" + self.visible_branches = ("default",) class SocialSecurityInputsMachine(RuleBasedStateMachine): - """All writes may slot earlier clones, branches, or restored copies.""" + """All writes may target earlier clones, branches, or restored copies.""" def __init__(self): super().__init__() @@ -115,7 +180,7 @@ def append(self, simulation, model): def write_key(model, component, period, data): # Reuse caller-known slots as well as drawing fresh ones. This makes # collisions between independent routes, copies and replay common. - keys = sorted(model.public.keys() | model.explicit.keys()) + keys = sorted(model.keys()) if ( data is not None and keys @@ -257,7 +322,9 @@ def dataset_load(self, slot, component, reported, period, amount, data=None): {self.name(component, reported): {period: [amount]}} ) simulation.build_from_dataset() - model = Model() + model = Model( + branch_name=model.branch_name, visible_branches=model.visible_branches + ) model.write(component, reported, period, amount, True, "dataset", initial=True) self.states[slot % len(self.states)] = (simulation, model) @@ -305,6 +372,7 @@ def builder_load( True, "yaml" if yaml else "builder", initial=True, + branch_name="default" if yaml else model.branch_name, ) if fresh: self.append(result, model) @@ -368,7 +436,13 @@ def copy_simulation(self, slot, branch): if branch else simulation.clone() ) - self.append(copied, copy.deepcopy(model)) + copied_model = copy.deepcopy(model) + if branch: + copied_model.branch_name = copied.branch_name + copied_model.visible_branches = tuple( + dict.fromkeys((copied.branch_name, *copied_model.visible_branches)) + ) + self.append(copied, copied_model) @rule(slot=TARGET) def dump_restore(self, slot): @@ -380,16 +454,25 @@ def dump_restore(self, slot): directory, self.backend.system ) copied_model = copy.deepcopy(model) - copied_model.dump_promotes_raw() + copied_model.flatten(promote_raw=True) self.append(restored, copied_model) - @rule(slot=TARGET) - def extract_reload(self, slot): + @rule(slot=TARGET, full_history=st.booleans()) + def extract_reload(self, slot, full_history=False): simulation, model = self.state(slot) - payload = simulation.extract_person() - restored = self.backend.Builder().build_from_dict(self.backend.system, payload) + if full_history: + payload = self.backend.hooks["export_source_inputs"](simulation) + restored = self.backend.hooks["reload_source_inputs"]( + payload, tax_benefit_system=self.backend.system + ) + else: + payload = simulation.extract_person() + restored = self.backend.Builder().build_from_dict( + self.backend.system, payload + ) copied_model = copy.deepcopy(model) copied_model.keep_registered() + copied_model.flatten() self.append(restored, copied_model) @rule(slot=TARGET, replay=st.booleans()) @@ -421,13 +504,25 @@ def apply(self, system): def reported_values_and_provenance_follow_intent(self): for simulation, model in self.states: assert simulation._user_input_contexts == [] - for key in model.public.keys() | model.explicit.keys(): + assert not hasattr(simulation, "_social_security_automatic_reported_inputs") + for key in model.keys(): component, period = key - expected = model.explicit.get(key, model.public.get(key)) + expected, expected_automatic = model.reported(component, period) reported = f"{component}_reported" actual = simulation.get_holder(reported).get_array( as_period(period), simulation.branch_name ) + if expected is None: + assert actual is None or simulation.get_holder(reported).is_derived( + as_period(period), simulation.branch_name + ), ( + key, + "public source shadowed explicit history", + ) + assert not self.backend.hooks["_is_automatic_reported_input"]( + simulation, reported, as_period(period) + ) + continue assert actual is not None, (key, expected.route, "missing twin") assert float(actual[0]) == expected.amount, ( key, @@ -438,7 +533,7 @@ def reported_values_and_provenance_follow_intent(self): automatic = self.backend.hooks["_is_automatic_reported_input"]( simulation, reported, as_period(period) ) - assert automatic == (key not in model.explicit), ( + assert automatic == expected_automatic, ( key, expected.route, "provenance", @@ -448,20 +543,23 @@ def reported_values_and_provenance_follow_intent(self): reported_holder = simulation.get_holder(f"{component}_reported") for period in reported_holder.get_input_periods(simulation.branch_name): key = (component, str(period)) - assert key in model.public or key in model.explicit, ( + assert model.reported(*key)[0] is not None, ( key, "unexplained reported input", ) holder = simulation.get_holder(component) for period in holder.get_input_periods(simulation.branch_name): key = (component, str(period)) - expected = model.overrides.get(key) + expected = model.visible(model.overrides, *key) assert expected is not None, (key, "unexpected public override") assert ( float(holder.get_array(period, simulation.branch_name)[0]) == expected.amount ) - for (component, period), expected in model.overrides.items(): + for component, period in model.keys(): + expected = model.visible(model.overrides, component, period) + if expected is None: + continue actual = simulation.get_holder(component).get_array( as_period(period), simulation.branch_name ) diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 3682e1b5cb8..76abfd25dc5 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -407,6 +407,18 @@ def set_input(self, variable, input_period, value): {"social_security_retirement": {2024: [24_000]}}, simulation=simulation, ) + # Combine the builder replacement with extraction before calculations + # add lookback caches to Core's first-known-period extraction path. + payload = simulation.extract_person(0) + reloaded = Simulation(situation=payload, tax_benefit_system=system) + result["dataset_builder_extraction_reload"] = _snapshot( + reloaded, (2024, 2026), INGRESS_OUTCOMES + ) + reloaded.set_input("social_security_retirement", 2024, [36_000]) + result["dataset_builder_extraction_rewrite"] = _snapshot( + reloaded, (2024, 2026), INGRESS_OUTCOMES + ) + del reloaded copied = simulation.clone(clone_tax_benefit_system=False) copied.set_input("social_security_retirement", 2026, [18_000]) SimulationBuilder().build_from_variables( @@ -837,6 +849,8 @@ def differential(tmp_path_factory): "direct_public_holder", "subclass_public_set_input_holder", "raw_core_dataset", + "dataset_builder_extraction_reload", + "dataset_builder_extraction_rewrite", "dataset_existing_builder_overwrite", "dataset_builder_clone_overwrite", "builder_failed_initial_loading", diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index 1202fb6a552..774d64617ea 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -648,15 +648,16 @@ def store_raw_source(holder, requested_period, array): assert amount(simulation, disability, 2026) == 0 -def test_raw_memory_mirror_keeps_a_surviving_disk_source_automatic(tmp_path): +def test_raw_memory_source_keeps_a_surviving_disk_input_available(tmp_path): public = SOCIAL_SECURITY_COMPONENTS[0] reported = f"{public}_reported" simulation = make_simulation() + source_holder = simulation.get_holder(public) reported_holder = simulation.get_holder(reported) - reported_holder._disk_storage = reported_holder.create_disk_storage( + source_holder._disk_storage = source_holder.create_disk_storage( tmp_path, preserve=True ) - reported_holder._on_disk_storable = True + source_holder._on_disk_storable = True simulation.memory_config = SimpleNamespace( max_memory_occupation_pc=0, priority_variables=system.variables, @@ -664,35 +665,78 @@ def test_raw_memory_mirror_keeps_a_surviving_disk_source_automatic(tmp_path): ) stored_period = period(2024) simulation.set_input(public, 2024, [12_000]) - assert reported_holder._memory_storage.get(stored_period) is None + assert source_holder._memory_storage.get(stored_period) is None np.testing.assert_array_equal( - reported_holder._disk_storage.get(stored_period), [12_000] + source_holder._disk_storage.get(stored_period), [12_000] ) + assert reported_holder._memory_storage.get(stored_period) is None - # A raw memory cache shadows the disk input without replacing its - # replayable tier or erasing that tier's automatic-source provenance. - reported_holder._on_disk_storable = False + # The read-time twin sees the raw memory source while registered-only + # exports retain the surviving disk input. No automatic twin is stored. + source_holder._on_disk_storable = False simulation.get_holder(public).put_in_cache(np.array([18_000]), stored_period) np.testing.assert_array_equal( - reported_holder._memory_storage.get(stored_period), [18_000] + source_holder._memory_storage.get(stored_period), [18_000] ) np.testing.assert_array_equal( simulation._get_recorded_input_array(reported, stored_period), [12_000] ) simulation._invalidate_all_caches() - assert reported_holder._memory_storage.get(stored_period) is None + assert source_holder._memory_storage.get(stored_period) is None np.testing.assert_array_equal( - reported_holder._disk_storage.get(stored_period), [12_000] + source_holder._disk_storage.get(stored_period), [12_000] ) - reported_holder._on_disk_storable = True + source_holder._on_disk_storable = True simulation.set_input(public, 2024, [24_000]) np.testing.assert_array_equal( - reported_holder._disk_storage.get(stored_period), [24_000] + source_holder._disk_storage.get(stored_period), [24_000] ) assert amount(simulation, reported, 2024) == 24_000 assert amount(simulation, public, 2026) == pytest.approx(27_044.32, abs=0.01) +@pytest.mark.parametrize("route", ("store_noninput", "cache")) +def test_dataset_source_survives_raw_twin_write_and_extraction(route): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + entities = ("household", "tax_unit", "spm_unit", "family", "marital_unit") + simulation = CoreSimulation( + dataset=pd.DataFrame( + { + "person_id": [1], + **{f"person_{entity}_id": [1] for entity in entities}, + **{f"{entity}_id": [1] for entity in entities}, + "age__2024": [70], + f"{public}__2024": [12_000], + } + ), + tax_benefit_system=system, + default_input_period=2024, + default_calculation_period=2024, + ) + holder = simulation.get_holder(reported) + assert not holder._memory_storage.get_known_branch_periods() + assert not any(name == reported for name, _, _ in simulation._user_input_keys) + if route == "cache": + holder.put_in_cache(np.array([6_000]), period(2024), derived=False) + else: + holder._set(period(2024), np.array([6_000]), is_input=False) + assert amount(simulation, reported, 2024) == 6_000 + assert simulation._get_recorded_input_array(reported, period(2024)) is None + + # Raw twins affect live calculations but are excluded from extraction. + # The independent registered public source must remain available on reload. + payload = json.loads(json.dumps(simulation.extract_person(0))) + reloaded = Simulation(situation=payload, tax_benefit_system=system) + assert amount(reloaded, reported, 2024) == 12_000 + assert amount(reloaded, public, 2024) == 12_000 + assert amount(reloaded, public, 2026) == pytest.approx(13_522.16, abs=0.01) + reloaded.set_input(public, 2024, [24_000]) + assert amount(reloaded, reported, 2024) == 24_000 + assert amount(reloaded, public, 2026) == pytest.approx(27_044.32, abs=0.01) + assert amount(simulation, reported, 2024) == 6_000 + + def test_derived_public_caches_never_become_reported_carry_sources(): retirement, disability = SOCIAL_SECURITY_COMPONENTS[:2] simulation = make_simulation() @@ -837,7 +881,7 @@ def test_complete_snapshot_preserves_histories_and_separate_final_overrides(): simulation.set_input(public, 2027, [0]) payload = json.loads(json.dumps(export_source_inputs(simulation))) person = payload["sources"]["people"]["person"] - assert set(person[f"{public}_reported"]) == {"2024", "2026", "2027"} + assert set(person[f"{public}_reported"]) == {"2024", "2026"} assert public not in person assert set(payload["overrides"][public]) == {"2026", "2027"} assert person["ss_aime_input"] == {"2024": 1_000, "2026": 2_000} @@ -850,9 +894,30 @@ def test_complete_snapshot_preserves_histories_and_separate_final_overrides(): assert input_periods(reloaded, public) == {"2026", "2027"} assert export_source_inputs(reloaded)["overrides"][public]["2027"] == [0] reloaded.set_input(public, 2027, [24_000]) - assert amount(reloaded, f"{public}_reported", 2027) == 24_000 - assert amount(reloaded, public, 2028) > 24_000 - assert amount(simulation, f"{public}_reported", 2027) == 0 + # An explicit twin at or before the request keeps priority over later + # public sources. The final public override remains independent of it. + source_index = system.parameters( + "2026-01-01" + ).calibration.gov.irs.soi.social_security_per_capita + expected_2027 = ( + 14_000 + * system.parameters( + "2027-01-01" + ).calibration.gov.irs.soi.social_security_per_capita + / source_index + ) + expected_2028 = ( + 14_000 + * system.parameters( + "2028-01-01" + ).calibration.gov.irs.soi.social_security_per_capita + / source_index + ) + assert amount(reloaded, f"{public}_reported", 2027) == pytest.approx(expected_2027) + assert amount(reloaded, public, 2028) == pytest.approx(expected_2028) + assert amount(simulation, f"{public}_reported", 2027) == pytest.approx( + expected_2027 + ) def test_legacy_extraction_uses_reported_sources_and_months_keep_flow_conversion(): diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 50c5aa54b18..5e855279737 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -1,21 +1,28 @@ -"""Compatibility ingress and complete, source-safe simulation snapshots. +"""Read-time Social Security sources without automatic reported inputs. -Initial ingress moves legacy benefit inputs to their pure reported sources. -A later public write overrides its period and refreshes its automatic source. -An explicitly supplied reported source retains precedence. +No automatic reported arrays or per-simulation origin maps are created. + +Initial public inputs remain supplied Core inputs on reserved source branches +of their public holders. Those branches are invisible to public calculations, +so benefit factors apply in the initial scored year. Later public inputs stay +normal final overrides. Reported holders expose a read-time source view; Core +still performs its own uprating, casting and derived-result caching. + +Explicit sources at or before a requested period take precedence. Snapshot +adapters serialize independent public source values and accept older origin +metadata for compatibility; they never serialize new sequence stamps. """ import json from contextvars import ContextVar +from dataclasses import dataclass from functools import wraps from pathlib import Path -from policyengine_core.holders.helpers import get_stored_input +from policyengine_core.data_storage.store_history import next_sequence_number from policyengine_core.periods import period as to_period from policyengine_core.simulations.simulation import NpEncoder -from policyengine_us.tools.period_branch import delete_input_values - SOCIAL_SECURITY_COMPONENTS = ( "social_security_retirement", @@ -24,548 +31,836 @@ "social_security_dependents", ) _REPORTED_COMPONENTS = tuple(f"{name}_reported" for name in SOCIAL_SECURITY_COMPONENTS) +_SOURCE_BRANCH = "__social_security_source__" +_SOURCE_DATA_FILE = "social_security_source_inputs.json" +_LEGACY_DUMP_FILE = "social_security_input_provenance.json" +# Keep the established envelope key so existing complete-snapshot callers +# need not change; version two contains actual source values, never stamps. +_PROVENANCE = "_social_security_input_provenance" - -_initial_social_security_ingress = ContextVar( - "initial_social_security_ingress", default=False -) -_initial_social_security_writes = ContextVar( - "initial_social_security_writes", default=None +_initial_social_security_ingress = ContextVar("initial_ss_ingress", default=False) +_initial_social_security_writes = ContextVar("initial_ss_writes", default=None) +_reported_request = ContextVar("reported_ss_request", default=None) +_subsampling_social_security_sources = ContextVar( + "subsampling_ss_sources", default=None ) -_reloaded_social_security_provenance = ContextVar( - "reloaded_social_security_provenance", default=None -) -_DUMP_PROVENANCE_FILE = "social_security_input_provenance.json" -_PROVENANCE = "_social_security_input_provenance" -_AUTOMATIC_INPUTS = "_social_security_automatic_reported_inputs" +_PHYSICAL = {} -def _recorded_input_location( - simulation, name, period, *, branch_names=None, include_unrecorded=False -): - """Locate a supplied input, or optionally a legacy nonderived cache seed. +@dataclass(frozen=True) +class _Value: + array: object + branch: str + tier: str + registered: bool + + +def _enabled(simulation): + return getattr( + getattr(simulation, "tax_benefit_system", None), + "_social_security_input_compatibility", + False, + ) + + +def _holder(simulation, name): + # Never create a holder while Core's dump traverses existing holders. + return simulation.persons._holders.get(name) + + +def _branches(simulation, branch_name=None): + if branch_name is None or branch_name == simulation.branch_name: + return simulation._get_visible_branch_names() + return simulation.get_holder(SOCIAL_SECURITY_COMPONENTS[0])._readable_branches( + branch_name + ) + + +def _source_branches(simulation, branch_name=None): + return tuple( + candidate + for branch in _branches(simulation, branch_name) + for candidate in (branch, f"{_SOURCE_BRANCH}{branch}") + ) + + +def _physical_periods(holder): + if holder is None: + return set() + return { + period + for storage in (holder._memory_storage, holder._disk_storage) + if storage is not None + for _, period in storage.get_known_branch_periods() + } + - Sequence stamps distinguish automatic sources from explicit replacements - even when the caller writes exactly the same amount. +def _physical_value(holder, period, branches, *, registered_only=False): + """Read supplied values and raw carry seeds, excluding formula caches. + + Core's per-tier input record decides whether an input survives. In + particular, nonderived memory seeds do not erase a recorded disk input + from a registered-only snapshot. """ - holder = simulation.get_holder(name) + if holder is None: + return None period = holder._storage_period(period) - input_keys = getattr(simulation, "_user_input_keys", ()) - if branch_names is None: - branch_names = simulation._get_visible_branch_names() - for branch_name in branch_names: - key = (name, branch_name, period) - for storage_name, storage, separator in ( - ("memory", holder._memory_storage, ":"), - ("disk", holder._disk_storage, "_"), + keys = getattr(holder.simulation, "_user_input_keys", ()) + for branch in branches: + for tier, storage in ( + ("memory", holder._memory_storage), + ("disk", holder._disk_storage), ): - if storage is None: + if storage is None or not storage.has(period, branch): continue - recorded = key in input_keys and holder._stores_user_input( - period, branch_name, storage_name - ) - if recorded or ( - include_unrecorded - and storage.has(period, branch_name) - and not storage.is_derived(period, branch_name) - ): - sequence = storage._sequence_numbers.get( - f"{branch_name}{separator}{period}" - ) - return key, (storage_name, sequence) + if storage.is_derived(period, branch): + continue + registered = ( + holder.variable.name, + branch, + period, + ) in keys and holder._stores_user_input(period, branch, tier) + if not registered_only or registered: + return _Value(storage.get(period, branch), branch, tier, registered) return None -def _is_automatic_reported_input(simulation, name, period, *, branch_names=None): - location = _recorded_input_location( - simulation, name, period, branch_names=branch_names, include_unrecorded=True +def _explicit_periods(simulation, name, branch_name=None, *, registered_only=False): + holder = _holder(simulation, name) + branches = _branches(simulation, branch_name) + return { + period + for period in _physical_periods(holder) + if _physical_value(holder, period, branches, registered_only=registered_only) + is not None + } + + +def _public_periods(simulation, name, branch_name=None, *, registered_only=False): + holder = _holder(simulation, name) + branches = _source_branches(simulation, branch_name) + return { + period + for period in _physical_periods(holder) + if _physical_value(holder, period, branches, registered_only=registered_only) + is not None + } + + +def _explicit_before(simulation, name, period, branch_name=None): + period = to_period(period) + # Core accepts an exact stored compound total before its annual period + # check, but cannot calculate or uprate another period into that total. + # An annual explicit source therefore cannot displace an exact multi-year + # public source. Annual requests retain Core's ordinary source ordering. + exact_only = period.unit == "year" and period.size > 1 + return any( + known == period if exact_only else known.start <= period.start + for known in _explicit_periods(simulation, name, branch_name) + ) + + +def _allow_fallback(simulation, name, period=None, branch_name=None): + request = _reported_request.get() + if request is not None and request[:2] == (simulation, name): + return request[2] + return period is None or not _explicit_before(simulation, name, period, branch_name) + + +def _effective_value( + simulation, name, period, branch_name=None, *, registered_only=False +): + explicit = _physical_value( + _holder(simulation, name), + period, + _branches(simulation, branch_name), + registered_only=registered_only, + ) + if explicit is not None: + return explicit, False + if not _allow_fallback(simulation, name, period, branch_name): + return None, False + public = name.removesuffix("_reported") + source = _physical_value( + _holder(simulation, public), + period, + _source_branches(simulation, branch_name), + registered_only=registered_only, ) - return location is not None and location[1] in getattr( - simulation, _AUTOMATIC_INPUTS, {} - ).get(location[0], ()) + return source, source is not None + + +def _is_automatic_reported_input(simulation, name, period, *, branch_names=None): + """Compatibility query computed from surviving sources, without stamps.""" + branch_name = None if branch_names is None else next(iter(branch_names), None) + _, fallback = _effective_value(simulation, name, period, branch_name) + return fallback -def _mark_automatic_reported_input(simulation, name, period, *, branch_name=None): - location = _recorded_input_location( - simulation, - name, +def _delete_exact(holder, period, branch_name, *, tier=None): + """Remove one storage slot, preserving compound-period neighbors. + + Core's general delete API also deletes contained periods. Initial routing + is a representation change and must not delete those other caller inputs. + """ + period = holder._storage_period(period) + if tier in (None, "memory"): + memory = holder._memory_storage + memory._arrays.pop(f"{branch_name}:{period}", None) + memory._unmark_dropped_keys() + memory._stop_sharing_dropped_keys() + memory._forget_dropped_numbers() + disk = holder._disk_storage + if disk is not None and tier in (None, "disk"): + disk._files.pop(f"{branch_name}_{period}", None) + disk._forget_dropped_keys() + recorded = getattr(holder, "_user_input_storage", None) + if recorded is not None: + recorded = dict(recorded) + slot = (branch_name, str(period)) + if tier is None: + recorded.pop(slot, None) + elif slot in recorded: + known_period, locations = recorded[slot] + recorded[slot] = (known_period, locations - {tier}) + holder._user_input_storage = recorded + if not holder._stores_user_input(period, branch_name): + holder.simulation._user_input_keys = { + key + for key in getattr(holder.simulation, "_user_input_keys", ()) + if key != (holder.variable.name, branch_name, period) + } + holder._evict_fast_cache(period, branch_name) + + +def _store_source(holder, period, array, branch_name, *, registered=True): + _PHYSICAL["set"]( + holder, period, - branch_names=None if branch_name is None else (branch_name,), - include_unrecorded=True, + array.copy(), + branch_name, + is_input=registered, + ) + + +def _refresh_input_variables(simulation): + existing = list(getattr(simulation, "input_variables", ())) + benefit_names = (*SOCIAL_SECURITY_COMPONENTS, *_REPORTED_COMPONENTS) + existing = [name for name in existing if name not in benefit_names] + visible = set(simulation._get_visible_branch_names()) + recorded_names = { + name + for name, branch, _ in getattr(simulation, "_user_input_keys", ()) + if branch in visible and name not in benefit_names + } + existing.extend( + name + for name in sorted(recorded_names) + if simulation._get_set_input_periods(name) ) - if location is not None: - # Core clones simulation attributes shallowly. Never mutate a shared - # provenance dictionary: each write belongs to this simulation alone. - automatic = dict(getattr(simulation, _AUTOMATIC_INPUTS, {})) - # A raw memory mirror can temporarily hide a surviving disk input. - # Keep each tier's origin so dropping that mirror reveals the original - # automatic disk source, rather than turning it into an explicit twin. - stamps = automatic.get(location[0], ()) - automatic[location[0]] = frozenset( - stamp for stamp in stamps if stamp[0] != location[1][0] - ) | {location[1]} - setattr(simulation, _AUTOMATIC_INPUTS, automatic) - - -def _automatic_reported_metadata(simulation, *, exported_periods=None): - automatic = {} - for public_name in SOCIAL_SECURITY_COMPONENTS: - name = f"{public_name}_reported" - # Dumping iterates this dictionary. Inspect existing holders without - # creating missing twins while Core is iterating over their values. - holder = simulation.persons._holders.get(name) + for name in benefit_names: + if simulation._get_set_input_periods(name): + existing.append(name) + simulation.input_variables = list(dict.fromkeys(existing)) + + +def migrate_social_security_inputs( + simulation, *, include_baseline=True, loaded_public_inputs=None +): + """Route successful initial public stores into invisible source branches. + + A caller-supplied same-period twin keeps the public value as a deliberate + final override. The independent public source is still preserved, so a + later raw twin's removal can reveal it again. Composed builders finalize + only stores belonging to that loading operation. + """ + moved = set() + for name in SOCIAL_SECURITY_COMPONENTS: + holder = _holder(simulation, name) if holder is None: continue - for period in holder.get_input_periods(simulation.branch_name): + reported = f"{name}_reported" + for branch in _branches(simulation): + source_branch = f"{_SOURCE_BRANCH}{branch}" + for period in tuple(_physical_periods(holder)): + value = _physical_value(holder, period, (branch,), registered_only=True) + if value is None: + continue + if loaded_public_inputs is not None: + if (name, branch, period) not in loaded_public_inputs: + continue + elif _physical_value(holder, period, (source_branch,)) is not None: + # A constructor's second migration must retain later + # public overrides and must not reclassify source copies. + continue + explicit = _physical_value( + _holder(simulation, reported), period, _branches(simulation) + ) + # Routing removes an exact public override and exposes its + # formula. A custom input handler may already have calculated + # values from that override while the builder was loading. + simulation._drop_values_that_may_depend_on(name, period) + _purge_derived_view(simulation, name) + _store_source(holder, period, value.array, source_branch) + if explicit is None: + _delete_exact(holder, period, branch) + moved.add(name) + _refresh_input_variables(simulation) + baseline = getattr(simulation, "baseline", None) + if include_baseline and baseline is not None and baseline is not simulation: + migrate_social_security_inputs(baseline, include_baseline=False) + return moved + + +def _automatic_reported_metadata( + simulation, *, exported_periods=None, person_indices=None, include_raw=False +): + """Serialize source data, not automatic provenance or sequence stamps. + + Physical public values are sufficient to reconstruct the read-time view. + ``projected_reported`` identifies only the compatibility projections in + the surrounding payload, so reload does not mistake them for new explicit + twins. It is derived at export and is never kept on the simulation. + """ + sources = [] + projected = {} + for public in SOCIAL_SECURITY_COMPONENTS: + holder = _holder(simulation, public) + if holder is None: + continue + for kind, branches in ( + ("initial", tuple(f"{_SOURCE_BRANCH}{b}" for b in _branches(simulation))), + ("override", tuple(_branches(simulation))), + ): + for period in _physical_periods(holder): + value = _physical_value( + holder, period, branches, registered_only=not include_raw + ) + if value is None: + continue + array = value.array + if person_indices is not None: + array = array[person_indices] + sources.append( + { + "name": public, + "period": str(period), + "kind": kind, + "values": array, + "registered": value.registered, + } + ) + reported = f"{public}_reported" + for period in _public_periods(simulation, public, registered_only=True): if ( exported_periods is not None - and (name, str(period)) not in exported_periods + and (reported, str(period)) not in exported_periods ): continue - if _is_automatic_reported_input(simulation, name, period): - automatic.setdefault(name, []).append(str(period)) - return {"version": 1, "automatic_reported": automatic} if automatic else None + value, fallback = _effective_value( + simulation, reported, period, registered_only=True + ) + if value is not None and fallback: + projected.setdefault(reported, []).append(str(period)) + if not sources and not projected: + return None + return json.loads( + json.dumps( + {"version": 2, "public_sources": sources, "projected_reported": projected}, + cls=NpEncoder, + ) + ) -def _restore_automatic_reported_metadata(simulation, metadata): +def _restore_automatic_reported_metadata(simulation, metadata, *, dumped=False): + """Accept old automatic projections; restore new independent public data.""" if metadata is None: return - if not isinstance(metadata, dict) or metadata.get("version") != 1: - raise ValueError("Unsupported Social Security input provenance version") - automatic = metadata.get("automatic_reported", {}) - if not isinstance(automatic, dict): - raise ValueError("Social Security input provenance must contain a mapping") - for name, periods in automatic.items(): - if name not in _REPORTED_COMPONENTS or not isinstance(periods, list): - raise ValueError("Invalid automatic Social Security reported input") - for period in periods: - _mark_automatic_reported_input(simulation, name, to_period(period)) + if not isinstance(metadata, dict) or metadata.get("version") not in (1, 2): + raise ValueError("Unsupported Social Security source-data version") + if metadata["version"] == 1: + # Older snapshots stored real automatic reported arrays. Move those + # values into independent public source slots and remove their copies. + for name, periods in metadata.get("automatic_reported", {}).items(): + if name not in _REPORTED_COMPONENTS or not isinstance(periods, list): + raise ValueError("Invalid legacy Social Security source data") + holder = _holder(simulation, name) + for text in periods: + period = to_period(text) + value = _physical_value(holder, period, _branches(simulation)) + if value is None: + continue + source = simulation.get_holder(name.removesuffix("_reported")) + _store_source( + source, + period, + value.array, + f"{_SOURCE_BRANCH}{simulation.branch_name}", + registered=value.registered, + ) + _delete_exact(holder, period, value.branch) + else: + for name, periods in metadata.get("projected_reported", {}).items(): + if name not in _REPORTED_COMPONENTS or not isinstance(periods, list): + raise ValueError("Invalid projected Social Security source") + holder = _holder(simulation, name) + if holder is not None: + for text in periods: + for branch in _branches(simulation): + _delete_exact(holder, to_period(text), branch) + import numpy as np + + for entry in metadata.get("public_sources", ()): + name, kind = entry.get("name"), entry.get("kind") + if name not in SOCIAL_SECURITY_COMPONENTS or kind not in ( + "initial", + "override", + ): + raise ValueError("Invalid Social Security public source") + branch = simulation.branch_name + if kind == "initial": + branch = f"{_SOURCE_BRANCH}{branch}" + _store_source( + simulation.get_holder(name), + to_period(entry["period"]), + np.asarray(entry["values"]), + branch, + registered=True if dumped else entry.get("registered", True), + ) + _refresh_input_variables(simulation) -def install_social_security_input_compatibility(): - """Extend Core's ingress routes only for systems opting into this policy. - - Core's entity builder writes holders directly, and its variable-level - ``set_input`` handler runs only for mismatched period units. Builder - and dataset adapters migrate initial sources. Holder-level interception - preserves subsequent public overrides from every supported input route: - simulation setters, direct holder writes, and custom period helpers. - Nonderived cache seeds retain Core's carryover semantics without becoming - registered inputs; calculated (derived) caches never synchronize. +def _select_subsample_sources(simulation, metadata, person_ids): + """Apply Core's retained person rows to independent public source data. + + Core keeps selected IDs and adjusts weights rather than duplicating rows. + Using those IDs also follows its entity rebuild when the original IDs are + nonconsecutive or the retained groups have different original sizes. """ + if metadata is None: + return None + positions = {person_id: index for index, person_id in enumerate(person_ids)} + retained = [positions[person_id] for person_id in simulation.persons.ids] + return { + **metadata, + "public_sources": [ + { + **entry, + "values": [entry["values"][index] for index in retained], + } + for entry in metadata.get("public_sources", ()) + ], + } + + +def _purge_derived_view(simulation, public_name): + """Discard only calculated view arrays; retain raw sources and inputs.""" + for name in (public_name, f"{public_name}_reported"): + holder = _holder(simulation, name) + if holder is None: + continue + for tier, storage in ( + ("memory", holder._memory_storage), + ("disk", holder._disk_storage), + ): + if storage is None: + continue + for branch, period in tuple(storage.get_known_branch_periods()): + if storage.is_derived(period, branch): + _delete_exact(holder, period, branch, tier=tier) + simulation._fast_cache = {} + + +def install_social_security_input_compatibility(): + from policyengine_core.holders import Holder from policyengine_core.simulations import Simulation as CoreSimulation from policyengine_core.simulations.simulation_builder import SimulationBuilder - from policyengine_core.holders import Holder + from policyengine_core.tools import simulation_dumper - def tracks_public_input(holder): - simulation = getattr(holder, "simulation", None) - return holder.variable.name in SOCIAL_SECURITY_COMPONENTS and getattr( - getattr(simulation, "tax_benefit_system", None), - "_social_security_input_compatibility", - False, + if getattr(Holder.get_array, "_social_security_readtime", False): + return + if getattr(Holder._set, "_social_security_compatibility", False): + raise RuntimeError( + "Replace eager compatibility before installing read-time hooks" ) - def synchronize_public_input(holder, period, branch_name, *, is_input): - simulation = holder.simulation - variable_name = holder.variable.name - period = holder._storage_period(period) - reported_name = f"{variable_name}_reported" - reported_holder = simulation.get_holder(reported_name) - readable_branches = reported_holder._readable_branches(branch_name) - reported_location = _recorded_input_location( - simulation, - reported_name, - period, - branch_names=readable_branches, - include_unrecorded=True, - ) - if reported_location is None or _is_automatic_reported_input( - simulation, reported_name, period, branch_names=readable_branches + _PHYSICAL.update( + set=Holder._set, + get_array=Holder.get_array, + get_known_periods=Holder.get_known_periods, + get_input_periods=Holder.get_input_periods, + recorded_periods=CoreSimulation._get_set_input_periods, + recorded_array=CoreSimulation._get_recorded_input_array, + calculate=CoreSimulation._calculate, + ) + + @wraps(_PHYSICAL["get_array"]) + def get_array(holder, period, branch_name="default"): + simulation, name = holder.simulation, holder.variable.name + if ( + name not in _REPORTED_COMPONENTS + or not _enabled(simulation) + or holder.variable.is_neutralized ): - supplied_value = ( - get_stored_input(holder, period, branch_name) - if is_input - else holder._get_array_from_storage(period, branch_name) - ) - if supplied_value is not None: - if is_input: - reported_holder.set_input( - period, supplied_value.copy(), branch_name - ) - else: - # Core also uprates from legacy, nonderived cache stores. - # Mirror their source semantics without promoting them to - # replayable/exported inputs. Calculated caches never enter. - reported_holder._set( - period, supplied_value.copy(), branch_name, is_input=False - ) - _mark_automatic_reported_input( - simulation, reported_name, period, branch_name=branch_name + return _PHYSICAL["get_array"](holder, period, branch_name) + value, fallback = _effective_value(simulation, name, period, branch_name) + if value is not None: + if fallback: + simulation._get_store_history().record_derived( + name.removesuffix("_reported"), next_sequence_number() ) - if not is_input: - return - # Core clones share this legacy list. Replace it privately so a - # later override is extracted without changing the source clone. - simulation.input_variables = list( - dict.fromkeys( - [ - *getattr(simulation, "input_variables", ()), - variable_name, - reported_name, + return value.array + return _PHYSICAL["get_array"](holder, period, branch_name) + + get_array._social_security_readtime = True + Holder.get_array = get_array + + @wraps(_PHYSICAL["get_known_periods"]) + def get_known_periods(holder): + values = _PHYSICAL["get_known_periods"](holder) + if holder.variable.name in _REPORTED_COMPONENTS and _enabled(holder.simulation): + if _allow_fallback(holder.simulation, holder.variable.name): + values = [ + *values, + *( + period + for period in _public_periods( + holder.simulation, + holder.variable.name.removesuffix("_reported"), + ) + if _effective_value( + holder.simulation, holder.variable.name, period + )[0] + is not None + ), ] + return list(dict.fromkeys(values)) + + Holder.get_known_periods = get_known_periods + + @wraps(_PHYSICAL["get_input_periods"]) + def get_input_periods(holder, branch_name="default"): + simulation, name = holder.simulation, holder.variable.name + if name not in _REPORTED_COMPONENTS or not _enabled(simulation): + return _PHYSICAL["get_input_periods"](holder, branch_name) + periods = _explicit_periods(simulation, name, branch_name) + if _allow_fallback(simulation, name, branch_name=branch_name): + periods |= _public_periods( + simulation, name.removesuffix("_reported"), branch_name ) + return [ + period + for period in periods + if _effective_value(simulation, name, period, branch_name)[0] is not None + ] + + Holder.get_input_periods = get_input_periods + + @wraps(_PHYSICAL["calculate"]) + def calculate(simulation, name, period=None): + if ( + name not in _REPORTED_COMPONENTS + or not _enabled(simulation) + or period is None + ): + return _PHYSICAL["calculate"](simulation, name, period) + token = _reported_request.set( + (simulation, name, not _explicit_before(simulation, name, period)) ) + try: + return _PHYSICAL["calculate"](simulation, name, period) + finally: + _reported_request.reset(token) + + CoreSimulation._calculate = calculate + + @wraps(_PHYSICAL["recorded_periods"]) + def recorded_periods(simulation, name): + periods = set(_PHYSICAL["recorded_periods"](simulation, name)) + if name in _REPORTED_COMPONENTS and _enabled(simulation): + periods |= _public_periods( + simulation, name.removesuffix("_reported"), registered_only=True + ) + periods = { + period + for period in periods + if _effective_value(simulation, name, period, registered_only=True)[0] + is not None + } + return sorted(periods, key=str) + + CoreSimulation._get_set_input_periods = recorded_periods - original_store = Holder._set - if not getattr(original_store, "_social_security_compatibility", False): + @wraps(_PHYSICAL["recorded_array"]) + def recorded_array(simulation, name, period): + if name not in _REPORTED_COMPONENTS or not _enabled(simulation): + return _PHYSICAL["recorded_array"](simulation, name, period) + value, _ = _effective_value(simulation, name, period, registered_only=True) + return None if value is None else value.array - @wraps(original_store) - def store_input( + CoreSimulation._get_recorded_input_array = recorded_array + + @wraps(_PHYSICAL["set"]) + def store( + holder, + period, + value, + branch_name="default", + validate_nan=False, + derived=False, + is_input=None, + sequence_number=None, + ): + simulation, name = holder.simulation, holder.variable.name + tracked = ( + not derived + and name in (*SOCIAL_SECURITY_COMPONENTS, *_REPORTED_COMPONENTS) + and _enabled(simulation) + ) + if tracked: + public = name.removesuffix("_reported") + simulation._drop_values_that_may_depend_on(public, period) + _purge_derived_view(simulation, public) + result = _PHYSICAL["set"]( holder, period, value, - branch_name="default", - validate_nan=False, - derived=False, - is_input=None, - sequence_number=None, - ): - track_source = not derived and tracks_public_input(holder) - # Mirror Core's classification and effective branch, including - # restore contexts and custom handlers that store before raising. - contexts = ( - getattr(holder.simulation, "_user_input_contexts", None) - if track_source - else None - ) - effective_input = bool(contexts) if is_input is None else is_input - stored_branch = ( + branch_name, + validate_nan=validate_nan, + derived=derived, + is_input=is_input, + sequence_number=sequence_number, + ) + if tracked: + contexts = getattr(simulation, "_user_input_contexts", ()) + registered = bool(contexts) if is_input is None else is_input + effective_branch = ( contexts[-1] - if effective_input and contexts and branch_name == "default" + if registered and contexts and branch_name == "default" else branch_name ) - result = original_store( - holder, - period, - value, - branch_name, - validate_nan=validate_nan, - derived=derived, - is_input=is_input, - sequence_number=sequence_number, - ) - loading_writes = _initial_social_security_writes.get() - if track_source and effective_input and loading_writes is not None: - loading_writes.setdefault(holder.simulation, {}).setdefault( - holder.variable.name, set() - ).add(holder._storage_period(period)) - # Initial registered inputs wait until every explicit twin is - # loaded. Raw cache seeds never enter that migration, so mirror - # them immediately without making either value a supplied input. - if track_source and ( - not _initial_social_security_ingress.get() or not effective_input - ): - synchronize_public_input( - holder, period, stored_branch, is_input=effective_input + writes = _initial_social_security_writes.get() + if name in SOCIAL_SECURITY_COMPONENTS and registered and writes is not None: + writes.setdefault(simulation, set()).add( + (name, effective_branch, holder._storage_period(period)) ) - return result - - store_input._social_security_compatibility = True - Holder._set = store_input - - original_dataset_build = CoreSimulation.build_from_dataset - if not getattr(original_dataset_build, "_social_security_compatibility", False): - - @wraps(original_dataset_build) - def build_from_dataset(simulation): - if not getattr( - simulation.tax_benefit_system, - "_social_security_input_compatibility", - False, - ): - return original_dataset_build(simulation) + elif name in SOCIAL_SECURITY_COMPONENTS: + # One logical public slot has one latest value. Replacing a + # source on this same branch retires its initial-routing + # copy; an inherited source on another branch remains a + # separate parent input, as ordinary Core fallback requires. + _delete_exact( + holder, + holder._storage_period(period), + f"{_SOURCE_BRANCH}{effective_branch}", + ) + _refresh_input_variables(simulation) + return result + + Holder._set = store + + def wrap_loading(original, *, dataset=False): + @wraps(original) + def load(first, *args, **kwargs): + system = ( + first.tax_benefit_system + if dataset + else args[0] + if args + else kwargs["tax_benefit_system"] + ) + if not getattr(system, "_social_security_input_compatibility", False): + return original(first, *args, **kwargs) + metadata = None + if not dataset: + input_dict = args[1] if len(args) > 1 else kwargs["input_dict"] + metadata = input_dict.get(_PROVENANCE) + if _PROVENANCE in input_dict: + input_dict = { + k: v for k, v in input_dict.items() if k != _PROVENANCE + } + if len(args) > 1: + args = (args[0], input_dict, *args[2:]) + else: + kwargs = {**kwargs, "input_dict": input_dict} outermost = not _initial_social_security_ingress.get() token = _initial_social_security_ingress.set(True) writes = _initial_social_security_writes.get() if writes is None: writes = {} writes_token = _initial_social_security_writes.set(writes) + result = None try: - result = original_dataset_build(simulation) + result = original(first, *args, **kwargs) + return result finally: _initial_social_security_writes.reset(writes_token) _initial_social_security_ingress.reset(token) if outermost: - # A successful store remains an input even when a later - # setter raises. Reconcile before propagating that error. - migrate_social_security_inputs( - simulation, loaded_public_inputs=writes.get(simulation, {}) - ) - reloaded = _reloaded_social_security_provenance.get() - if reloaded is not None and reloaded[0] is simulation: - _restore_automatic_reported_metadata(simulation, reloaded[1]) - return result + targets = dict(writes) + target = first if dataset else result + if target is None and not dataset: + target = kwargs.get("simulation") + if target is None and len(args) > 2: + target = args[2] + if target is not None: + targets.setdefault(target, set()) + for target, loaded in targets.items(): + # During routing, projected twins temporarily retain + # any both-name public final overrides. Conversion + # then replaces those projection arrays with the + # independently serialized public source data. + migrate_social_security_inputs( + target, loaded_public_inputs=loaded + ) + rebuilding = _subsampling_social_security_sources.get() + if ( + dataset + and rebuilding is not None + and rebuilding[0] is target + ): + metadata = _select_subsample_sources( + target, rebuilding[1], rebuilding[2] + ) + _restore_automatic_reported_metadata(target, metadata) - build_from_dataset._social_security_compatibility = True - CoreSimulation.build_from_dataset = build_from_dataset + return load + + CoreSimulation.build_from_dataset = wrap_loading( + CoreSimulation.build_from_dataset, dataset=True + ) + for name in ("build_from_dict", "build_from_entities", "build_from_variables"): + setattr(SimulationBuilder, name, wrap_loading(getattr(SimulationBuilder, name))) original_subsample = CoreSimulation.subsample - if not getattr(original_subsample, "_social_security_compatibility", False): - - @wraps(original_subsample) - def subsample(simulation, *args, **kwargs): - if not getattr( - simulation.tax_benefit_system, - "_social_security_input_compatibility", - False, - ): - return original_subsample(simulation, *args, **kwargs) - # Restore stamps during dataset rebuilding, before Core replays - # inputs and clones the rebuilt reform baseline. - token = _reloaded_social_security_provenance.set( - (simulation, _automatic_reported_metadata(simulation)) + + @wraps(original_subsample) + def subsample(simulation, *args, **kwargs): + if not _enabled(simulation): + return original_subsample(simulation, *args, **kwargs) + token = _subsampling_social_security_sources.set( + ( + simulation, + _automatic_reported_metadata(simulation), + tuple(simulation.persons.ids), ) - try: - return original_subsample(simulation, *args, **kwargs) - finally: - _reloaded_social_security_provenance.reset(token) + ) + try: + return original_subsample(simulation, *args, **kwargs) + finally: + _subsampling_social_security_sources.reset(token) - subsample._social_security_compatibility = True - CoreSimulation.subsample = subsample + CoreSimulation.subsample = subsample - from policyengine_core.tools import simulation_dumper + original_extract = CoreSimulation.extract_person - original_dump_holder = simulation_dumper._dump_holder - if not getattr(original_dump_holder, "_social_security_compatibility", False): - - @wraps(original_dump_holder) - def dump_holder(holder, directory): - result = original_dump_holder(holder, directory) - simulation = holder.simulation - name = holder.variable.name - if name in _REPORTED_COMPONENTS and getattr( - simulation.tax_benefit_system, - "_social_security_input_compatibility", - False, - ): - exported_periods = { - (name, str(period)) for period in holder.get_known_periods() + @wraps(original_extract) + def extract_person(simulation, index=0, exclude_entities=("state",)): + payload = original_extract(simulation, index, exclude_entities) + if not _enabled(simulation): + return payload + people = payload.get(simulation.persons.entity.plural, {}) + indices = [] + exported = set() + for person_name, person in people.items(): + person_index = ( + int(person_name.removeprefix(f"{simulation.persons.entity.key}_")) - 1 + ) + indices.append(person_index) + for name in (*SOCIAL_SECURITY_COMPONENTS, *_REPORTED_COMPONENTS): + values = { + str(period): float(array[person_index]) + for period in simulation._get_set_input_periods(name) + if (array := simulation._get_recorded_input_array(name, period)) + is not None } - metadata = _automatic_reported_metadata( - simulation, exported_periods=exported_periods - ) - if metadata is not None: - path = Path(directory) / name / _DUMP_PROVENANCE_FILE - path.write_text(json.dumps(metadata)) - return result + if values: + person[name] = values + exported.update((name, period) for period in values) + else: + person.pop(name, None) + metadata = _automatic_reported_metadata( + simulation, exported_periods=exported, person_indices=indices + ) + if metadata is not None: + payload[_PROVENANCE] = metadata + return payload - dump_holder._social_security_compatibility = True - simulation_dumper._dump_holder = dump_holder + CoreSimulation.extract_person = extract_person - original_restore_holder = simulation_dumper._restore_holder - if not getattr(original_restore_holder, "_social_security_compatibility", False): + original_dump = simulation_dumper._dump_holder - @wraps(original_restore_holder) - def restore_holder( - simulation, variable, directory, inputs=False, calculated_number=None + @wraps(original_dump) + def dump_holder(holder, directory): + result = original_dump(holder, directory) + if not _enabled(holder.simulation): + return result + if holder.variable.name in (*SOCIAL_SECURITY_COMPONENTS, *_REPORTED_COMPONENTS): + # A virtual source need not have a physically visible array. + # Core creates the directory before discovering that; leave no + # empty holder folder for its restore parser to misinterpret. + storage_path = Path(directory) / holder.variable.name + if storage_path.exists() and not any(storage_path.iterdir()): + storage_path.rmdir() + if holder.variable.name in SOCIAL_SECURITY_COMPONENTS: + metadata = _automatic_reported_metadata(holder.simulation, include_raw=True) + if metadata is not None: + # Core enumerates every root entry as a variable to restore. + # Keep source data inside a valid holder directory, including + # when all this holder's actual inputs use hidden branches. + storage_path = Path(directory) / holder.variable.name + storage_path.mkdir(exist_ok=True) + (storage_path / _SOURCE_DATA_FILE).write_text(json.dumps(metadata)) + return result + + simulation_dumper._dump_holder = dump_holder + original_restore = simulation_dumper._restore_holder + + @wraps(original_restore) + def restore_holder( + simulation, variable, directory, inputs=False, calculated_number=None + ): + if not _enabled(simulation): + return original_restore( + simulation, variable, directory, inputs, calculated_number + ) + if not inputs and not getattr( + simulation, "_social_security_source_data_restored", False ): - if not getattr( - simulation.tax_benefit_system, - "_social_security_input_compatibility", - False, - ): - return original_restore_holder( - simulation, variable, directory, inputs, calculated_number + source_path = next( + ( + path + for name in SOCIAL_SECURITY_COMPONENTS + if (path := Path(directory) / name / _SOURCE_DATA_FILE).exists() + ), + None, + ) + if source_path is not None: + _restore_automatic_reported_metadata( + simulation, json.loads(source_path.read_text()), dumped=True ) - if not inputs and not getattr( - simulation, "_social_security_dump_inputs_restored", False - ): - # Core finishes the input pass before restoring any calculated - # arrays. Load all twins first, so directory order cannot turn - # an explicit twin into an automatic source or vice versa. - migrate_social_security_inputs(simulation) + else: for name in _REPORTED_COMPONENTS: - path = Path(directory) / name / _DUMP_PROVENANCE_FILE + path = Path(directory) / name / _LEGACY_DUMP_FILE if path.exists(): _restore_automatic_reported_metadata( - simulation, json.loads(path.read_text()) + simulation, json.loads(path.read_text()), dumped=True ) - simulation._social_security_dump_inputs_restored = True - token = _initial_social_security_ingress.set(True) - try: - return original_restore_holder( - simulation, variable, directory, inputs, calculated_number - ) - finally: - _initial_social_security_ingress.reset(token) - - restore_holder._social_security_compatibility = True - simulation_dumper._restore_holder = restore_holder - - original_extract = CoreSimulation.extract_person - if not getattr(original_extract, "_social_security_compatibility", False): - - @wraps(original_extract) - def extract_person(simulation, index=0, exclude_entities=("state",)): - payload = original_extract(simulation, index, exclude_entities) - if not getattr( - simulation.tax_benefit_system, - "_social_security_input_compatibility", - False, - ): - return payload - people = payload.get(simulation.persons.entity.plural, {}) - benefit_inputs = { - name: { - str(period): value - for period in simulation._get_set_input_periods(name) - if (value := simulation._get_recorded_input_array(name, period)) - is not None - } - for name in (*SOCIAL_SECURITY_COMPONENTS, *_REPORTED_COMPONENTS) - } - for person_name, person in people.items(): - # Core names extracted people by their original zero-based - # index plus one, even when extracting a later household. - person_index = ( - int(person_name.removeprefix(f"{simulation.persons.entity.key}_")) - - 1 - ) - for name, values in benefit_inputs.items(): - # Core exports its first known period, which may be a - # formula cache rather than the later public override. - # Keep every actual benefit input and its source history - # together, without converting calculated values to inputs. - if values: - person[name] = { - period: float(value[person_index]) - for period, value in values.items() - } - else: - person.pop(name, None) - exported_periods = { - (name, str(period)) - for person in people.values() - for name, values in person.items() - if name in _REPORTED_COMPONENTS - for period in values - } - metadata = _automatic_reported_metadata( - simulation, exported_periods=exported_periods + migrate_social_security_inputs(simulation) + simulation._social_security_source_data_restored = True + token = _initial_social_security_ingress.set(True) + try: + return original_restore( + simulation, variable, directory, inputs, calculated_number ) - if metadata is not None: - payload[_PROVENANCE] = metadata - return payload - - extract_person._social_security_compatibility = True - CoreSimulation.extract_person = extract_person - - original_invalidate = CoreSimulation._invalidate_all_caches - if not getattr(original_invalidate, "_social_security_compatibility", False): - - @wraps(original_invalidate) - def invalidate_all_caches(simulation): - automatic = [] - if getattr( - simulation.tax_benefit_system, - "_social_security_input_compatibility", - False, - ): - for key, stamps in getattr(simulation, _AUTOMATIC_INPUTS, {}).items(): - name, branch_name, period = key - location = _recorded_input_location( - simulation, name, period, branch_names=(branch_name,) - ) - if location is not None and location[1] in stamps: - automatic.append(key) - result = original_invalidate(simulation) - # Core replay assigns new store sequence numbers. Replaying an - # automatic source does not turn it into an explicit caller input. - for name, branch_name, period in automatic: - _mark_automatic_reported_input( - simulation, name, period, branch_name=branch_name - ) - return result + finally: + _initial_social_security_ingress.reset(token) - invalidate_all_caches._social_security_compatibility = True - CoreSimulation._invalidate_all_caches = invalidate_all_caches - - def wrap_builder(original_build): - @wraps(original_build) - def build(builder, tax_benefit_system, input_dict, simulation=None): - if not getattr( - tax_benefit_system, "_social_security_input_compatibility", False - ): - return original_build( - builder, tax_benefit_system, input_dict, simulation - ) - metadata = input_dict.get(_PROVENANCE) - if _PROVENANCE in input_dict: - input_dict = { - name: values - for name, values in input_dict.items() - if name != _PROVENANCE - } - outermost = not _initial_social_security_ingress.get() - token = _initial_social_security_ingress.set(True) - # Nested dict/entity/variable builders form one loading operation. - # Record actual stores after Core has handled their periods. - writes = _initial_social_security_writes.get() - if writes is None: - writes = {} - writes_token = _initial_social_security_writes.set(writes) - result = None - try: - result = original_build( - builder, tax_benefit_system, input_dict, simulation - ) - finally: - _initial_social_security_writes.reset(writes_token) - _initial_social_security_ingress.reset(token) - if outermost: - targets = dict(writes) - target = result if result is not None else simulation - if target is not None: - targets.setdefault(target, {}) - for target, loaded in targets.items(): - # Successful stores identify even a fresh build that - # raises before it can return its simulation. - visible_branches = target._get_visible_branch_names() - recorded_names = { - name - for name, branch_name, _ in target._user_input_keys - if branch_name in visible_branches - } - target.input_variables = list( - dict.fromkeys( - [ - *getattr(target, "input_variables", []), - *sorted( - name - for name in recorded_names - if target._get_set_input_periods(name) - ), - ] - ) - ) - migrate_social_security_inputs( - target, loaded_public_inputs=loaded - ) - _restore_automatic_reported_metadata(target, metadata) - return result - - build._social_security_compatibility = True - return build - - for method_name in ( - "build_from_dict", - "build_from_entities", - "build_from_variables", - ): - original_build = getattr(SimulationBuilder, method_name) - if not getattr(original_build, "_social_security_compatibility", False): - setattr(SimulationBuilder, method_name, wrap_builder(original_build)) + simulation_dumper._restore_holder = restore_holder + # Core's physical replay/drop/clone methods require no automatic-stamp + # adapter: the independent public source is a genuine registered input. + # The complete snapshot entry points below keep their established shape. from policyengine_us.tools.social_security_periods import ( install_social_security_period_compatibility, ) @@ -573,110 +868,8 @@ def build(builder, tax_benefit_system, input_dict, simulation=None): install_social_security_period_compatibility() -def migrate_social_security_inputs( - simulation, *, include_baseline=True, loaded_public_inputs=None -): - """Move legacy component inputs, period by period, into pure sources. - - An explicit twin in the same visible period keeps its priority. Automatic - twins refresh for public inputs stored by this loading operation. Earlier - public overrides are not initial inputs merely because another builder is - run on the same simulation. Without a loading record, migrate only legacy - public inputs with no twin. Recorded-input helpers exclude formula caches - and honor visible branch precedence. - """ - visible_branches = simulation._get_visible_branch_names() - migrated = set() - for public_name in SOCIAL_SECURITY_COMPONENTS: - reported_name = f"{public_name}_reported" - holder = simulation.get_holder(public_name) - # Snapshot before deleting: deleting a multi-year period also deletes - # contained supplied periods, including a shorter legacy source whose - # lexical sort order happens to place it after the longer period. - sources = {} - for period in simulation._get_set_input_periods(public_name): - location = _recorded_input_location( - simulation, reported_name, period, include_unrecorded=True - ) - if location is not None and ( - not _is_automatic_reported_input(simulation, reported_name, period) - or loaded_public_inputs is None - or period not in loaded_public_inputs.get(public_name, ()) - ): - continue - value = simulation._get_recorded_input_array(public_name, period) - if value is not None: - sources[period] = value.copy() - if not sources: - continue - retained_overrides = [] - for name, branch_name, period in tuple(simulation._user_input_keys): - if ( - name != public_name - or branch_name not in visible_branches - or period in sources - ): - continue - for storage_name, storage in ( - ("memory", holder._memory_storage), - ("disk", holder._disk_storage), - ): - if holder._stores_user_input(period, branch_name, storage_name): - retained_overrides.append( - (branch_name, period, storage.get(period, branch_name).copy()) - ) - break - for period, value in sources.items(): - simulation.set_input(reported_name, period, value.copy()) - _mark_automatic_reported_input(simulation, reported_name, period) - for period in sources: - for branch_name in visible_branches: - delete_input_values(holder, period, branch_name) - # Keep exact-period both-name overrides even inside a migrated longer - # source period. Restore each original branch slot through set_input, - # preserving its input classification and ordinary replay bookkeeping. - for branch_name, period, value in retained_overrides: - if not holder._stores_user_input(period, branch_name): - holder.set_input(period, value, branch_name) - migrated.add(public_name) - - # Core's legacy extraction reads this runtime list, which was assembled - # before migration. Retain public names only where actual overrides remain. - runtime_inputs = [] - for name in getattr(simulation, "input_variables", []): - if name not in migrated or simulation._get_set_input_periods(name): - runtime_inputs.append(name) - if name in migrated: - runtime_inputs.append(f"{name}_reported") - for public_name in SOCIAL_SECURITY_COMPONENTS: - reported_name = f"{public_name}_reported" - if simulation._get_set_input_periods(reported_name): - runtime_inputs.append(reported_name) - simulation.input_variables = list(dict.fromkeys(runtime_inputs)) - # Core creates this clone during super().__init__, before either PE-US - # wrapper reaches migration. Its copied public sources otherwise override - # formulas in the supplied year and disappear from the uprating route in - # later years. Each branch needs its own source/input-key migration. - baseline = getattr(simulation, "baseline", None) - if include_baseline and baseline is not None and baseline is not simulation: - migrate_social_security_inputs(baseline, include_baseline=False) - return migrated - - def export_source_inputs(simulation): - """Export all supplied periods, separating pure sources and final overrides. - - The JSON-compatible payload preserves entity memberships and every recorded - pure-input period, including selectors, dates, worker links and earnings - histories. Supplied formula-variable inputs have their own ``overrides`` - payload; calculated results never become sources. Only the current branch - and its visible ancestors are exported, using their normal precedence. - - This is the complete-history alternative to Core's legacy ``extract_person`` - API, which retains all supplied periods only for benefit components and their - reported twins. Export the desired small simulation or household; this - helper deliberately does not materialize a population dataset. - """ + """Export all registered source periods and final overrides.""" names = { key: [str(name) for name in population.ids] for key, population in simulation.populations.items() @@ -686,30 +879,23 @@ def export_source_inputs(simulation): entity = population.entity sources[entity.plural] = {name: {} for name in names[key]} if not entity.is_person: - # All PE-US group entities use the unbounded member role. for person_index, group_index in enumerate(population.members_entity_id): group = sources[entity.plural][names[key][int(group_index)]] group.setdefault("members", []).append(names["person"][person_index]) - overrides = {} for name, variable in simulation.tax_benefit_system.variables.items(): - periods = simulation._get_set_input_periods(name) - if not periods: - continue - for period in periods: + for period in simulation._get_set_input_periods(name): values = simulation._get_recorded_input_array(name, period) if values is None: continue if hasattr(values, "decode_to_str"): values = values.decode_to_str() if variable.is_input_variable(): - entity = variable.entity - for index, entity_name in enumerate(names[entity.key]): - source = sources[entity.plural][entity_name] + for index, entity_name in enumerate(names[variable.entity.key]): + source = sources[variable.entity.plural][entity_name] source.setdefault(name, {})[str(period)] = values[index] else: overrides.setdefault(name, {})[str(period)] = values - return json.loads( json.dumps( { @@ -728,12 +914,7 @@ def export_source_inputs(simulation): def reload_source_inputs(payload, **simulation_kwargs): - """Reload a complete source snapshot, then restore its final overrides. - - Pass the desired ``tax_benefit_system`` or ``reform`` through keyword - arguments. Loading overrides after construction preserves a public-only - final override instead of treating it as a legacy reported source. - """ + """Reload complete snapshots, including legacy automatic projections.""" from policyengine_us import Simulation if payload.get("version") != 1: @@ -743,15 +924,10 @@ def reload_source_inputs(payload, **simulation_kwargs): metadata = payload.get(_PROVENANCE) sources = payload["sources"] if metadata is not None: - # Restore before Core creates a reform baseline, so that clone also - # knows which twins came from legacy public inputs. sources = {**sources, _PROVENANCE: metadata} simulation = Simulation(situation=sources, **simulation_kwargs) - _restore_automatic_reported_metadata(simulation, metadata) for name, periods in payload["overrides"].items(): for period, values in periods.items(): simulation.set_input(name, period, values) - simulation.input_variables = list( - dict.fromkeys([*simulation.input_variables, *payload["overrides"]]) - ) + _refresh_input_variables(simulation) return simulation From 8c4640ffa925e369282856923d96401dd5e1d94c Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 13:36:33 -0700 Subject: [PATCH 38/76] Preserve benefit sources through branch and cache cleanup --- .../test_social_security_current_law.py | 3 ++ .../test_social_security_inputs.py | 2 + policyengine_us/tools/period_branch.py | 20 +++++++++ .../tools/social_security_inputs.py | 45 ++++++++++++++++++- 4 files changed, 68 insertions(+), 2 deletions(-) diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 76abfd25dc5..000f4470656 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -402,6 +402,9 @@ def set_input(self, variable, input_period, value): default_input_period=2024, default_calculation_period=2024, ) + # Existing builders must also replace a source after its scored-year + # formula (including its base) has already been cached. + simulation.calculate("social_security_retirement", 2024) SimulationBuilder().build_from_variables( system, {"social_security_retirement": {2024: [24_000]}}, diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index 774d64617ea..4920b7d8115 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -299,6 +299,7 @@ def test_existing_builder_refreshes_loaded_sources_and_keeps_earlier_overrides( inputs = {name: {2024: 12_000} for name in SOCIAL_SECURITY_COMPONENTS} simulation = make_core_simulation(inputs) for name in SOCIAL_SECURITY_COMPONENTS: + assert amount(simulation, name, 2024) == 12_000 simulation.set_input(name, 2025, [18_000]) if explicit_source is not None: simulation.set_input(f"{name}_reported", 2024, [explicit_source]) @@ -308,6 +309,7 @@ def test_existing_builder_refreshes_loaded_sources_and_keeps_earlier_overrides( simulation=simulation, ) for name in SOCIAL_SECURITY_COMPONENTS: + assert amount(simulation, name, 2024) == 24_000 assert amount(simulation, f"{name}_reported", 2024) == ( 24_000 if explicit_source is None else explicit_source ) diff --git a/policyengine_us/tools/period_branch.py b/policyengine_us/tools/period_branch.py index dd70ee44056..93d598229bb 100644 --- a/policyengine_us/tools/period_branch.py +++ b/policyengine_us/tools/period_branch.py @@ -65,10 +65,30 @@ def _input_keys(branch: Simulation) -> Set[Tuple[str, str, Period]]: keys and its ancestors'. """ visible_branches = set(branch._get_visible_branch_names()) + source_branches = set() + source_names = () + if getattr( + branch.tax_benefit_system, "_social_security_input_compatibility", False + ): + # Initial public benefit sources are genuine Core inputs stored on + # reserved branches, so public formulas can apply same-year levers. + # Keep them when clearing a branch's inherited formula results. + from policyengine_us.tools.social_security_inputs import ( + SOCIAL_SECURITY_COMPONENTS, + _SOURCE_BRANCH, + ) + + source_names = SOCIAL_SECURITY_COMPONENTS + source_branches = {f"{_SOURCE_BRANCH}{name}" for name in visible_branches} return { key for key in getattr(branch, "_user_input_keys", set()) if key[1] in visible_branches + or ( + key[0] in source_names + and key[1] in source_branches + and branch.get_holder(key[0])._stores_user_input(key[2], key[1]) + ) } diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 5e855279737..47a9295963c 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -275,6 +275,7 @@ def migrate_social_security_inputs( only stores belonging to that loading operation. """ moved = set() + purged = False for name in SOCIAL_SECURITY_COMPONENTS: holder = _holder(simulation, name) if holder is None: @@ -300,10 +301,30 @@ def migrate_social_security_inputs( # formula. A custom input handler may already have calculated # values from that override while the builder was loading. simulation._drop_values_that_may_depend_on(name, period) - _purge_derived_view(simulation, name) + if not purged: + # Core's dependency drops apply only on branches. A + # builder can replace a source on an ordinary simulation + # whose base or downstream values are already cached. + # Keep raw carry seeds and every supplied input intact. + _purge_derived_arrays(simulation) + purged = True _store_source(holder, period, value.array, source_branch) if explicit is None: - _delete_exact(holder, period, branch) + # A local source must also displace inherited public + # final overrides in this simulation's private snapshot. + # Retain each ancestor's independent public source for + # ordinary branch fallback after a later local raw write. + for visible in _branches(simulation): + inherited = _physical_value(holder, period, (visible,)) + if inherited is not None and visible != branch: + _store_source( + holder, + period, + inherited.array, + f"{_SOURCE_BRANCH}{visible}", + registered=inherited.registered, + ) + _delete_exact(holder, period, visible) moved.add(name) _refresh_input_variables(simulation) baseline = getattr(simulation, "baseline", None) @@ -471,6 +492,26 @@ def _purge_derived_view(simulation, public_name): simulation._fast_cache = {} +def _purge_derived_arrays(simulation): + """Discard formula results after initial routing changes their source. + + Core's general cache drop also removes nonregistered, nonderived carry + seeds. Those are independent caller values and remain valid here. + """ + for population in simulation.populations.values(): + for holder in population._holders.values(): + for tier, storage in ( + ("memory", holder._memory_storage), + ("disk", holder._disk_storage), + ): + if storage is None: + continue + for branch, period in tuple(storage.get_known_branch_periods()): + if storage.is_derived(period, branch): + _delete_exact(holder, period, branch, tier=tier) + simulation._fast_cache = {} + + def install_social_security_input_compatibility(): from policyengine_core.holders import Holder from policyengine_core.simulations import Simulation as CoreSimulation From b338d0a13ff7ed96d71ad4e790a5b3548da62f0c Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 13:44:48 -0700 Subject: [PATCH 39/76] Skip source routing before populations exist --- policyengine_us/tools/social_security_inputs.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 47a9295963c..80571556cf2 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -65,7 +65,8 @@ def _enabled(simulation): def _holder(simulation, name): # Never create a holder while Core's dump traverses existing holders. - return simulation.persons._holders.get(name) + persons = getattr(simulation, "persons", None) + return getattr(persons, "_holders", {}).get(name) def _branches(simulation, branch_name=None): @@ -275,6 +276,10 @@ def migrate_social_security_inputs( only stores belonging to that loading operation. """ moved = set() + # A loader can fail before Core creates its populations. There are then + # no successful holder writes to route, including in constructor probes. + if getattr(simulation, "persons", None) is None: + return moved purged = False for name in SOCIAL_SECURITY_COMPONENTS: holder = _holder(simulation, name) From 607106fe917f6b741019dd062b67b4ff771e3714 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 13:56:58 -0700 Subject: [PATCH 40/76] Supply group role metadata to the Core test backend --- .../tests/core/_social_security_state_backend.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/policyengine_us/tests/core/_social_security_state_backend.py b/policyengine_us/tests/core/_social_security_state_backend.py index bd055237505..e1c8d6a2a24 100644 --- a/policyengine_us/tests/core/_social_security_state_backend.py +++ b/policyengine_us/tests/core/_social_security_state_backend.py @@ -318,8 +318,15 @@ class PreservedInput: storage: str person = types.SimpleNamespace(key="person", plural="people", is_person=True) + member = types.SimpleNamespace( + key="member", plural="members", max=None, subroles=None + ) household = types.SimpleNamespace( - key="household", plural="households", is_person=False + key="household", + plural="households", + is_person=False, + roles=(member,), + flattened_roles=(member,), ) class Holder: From ce2ad34c0f1e455a9b7231cfa361b004ec98d8a3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 16:39:46 -0700 Subject: [PATCH 41/76] Refresh derived benefits after source writes Clear calculated bases and downstream arrays when a component source changes, retaining supplied inputs and raw seeds. Keep the monetary regression checks and prime the aggregate result too. Replace the physical-twin assertion with a default-branch value assertion: the read-time design deliberately stores no automatic twin. --- .../test_social_security_inputs.py | 15 ++++++++--- .../tools/social_security_inputs.py | 25 ++++--------------- 2 files changed, 17 insertions(+), 23 deletions(-) diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/microsimulation/test_social_security_inputs.py index 4920b7d8115..3f5328e40d5 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_inputs.py +++ b/policyengine_us/tests/microsimulation/test_social_security_inputs.py @@ -522,9 +522,12 @@ def test_explicit_noninput_store_keeps_its_branch_inside_an_input_context(): finally: simulation._user_input_contexts.pop() for name in (public, reported): - assert set(simulation.get_holder(name).get_known_branch_periods()) == { - ("default", period(2024)) - } + # The reported holder exposes the public source at read time; it + # need not store a separate twin on the default branch. + np.testing.assert_array_equal( + simulation.get_holder(name).get_array(period(2024), "default"), + [12_000], + ) assert input_periods(simulation, name) == set() assert amount(simulation, name, 2024) == 12_000 assert amount(simulation, public, 2026) == pytest.approx(13_522.16, abs=0.01) @@ -733,9 +736,15 @@ def test_dataset_source_survives_raw_twin_write_and_extraction(route): assert amount(reloaded, reported, 2024) == 12_000 assert amount(reloaded, public, 2024) == 12_000 assert amount(reloaded, public, 2026) == pytest.approx(13_522.16, abs=0.01) + assert amount(reloaded, "social_security", 2026) == pytest.approx( + 13_522.16, abs=0.01 + ) reloaded.set_input(public, 2024, [24_000]) assert amount(reloaded, reported, 2024) == 24_000 assert amount(reloaded, public, 2026) == pytest.approx(27_044.32, abs=0.01) + assert amount(reloaded, "social_security", 2026) == pytest.approx( + 27_044.32, abs=0.01 + ) assert amount(simulation, reported, 2024) == 6_000 diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 80571556cf2..c319a8cc5eb 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -479,26 +479,8 @@ def _select_subsample_sources(simulation, metadata, person_ids): } -def _purge_derived_view(simulation, public_name): - """Discard only calculated view arrays; retain raw sources and inputs.""" - for name in (public_name, f"{public_name}_reported"): - holder = _holder(simulation, name) - if holder is None: - continue - for tier, storage in ( - ("memory", holder._memory_storage), - ("disk", holder._disk_storage), - ): - if storage is None: - continue - for branch, period in tuple(storage.get_known_branch_periods()): - if storage.is_derived(period, branch): - _delete_exact(holder, period, branch, tier=tier) - simulation._fast_cache = {} - - def _purge_derived_arrays(simulation): - """Discard formula results after initial routing changes their source. + """Discard formula results after benefit sources change. Core's general cache drop also removes nonregistered, nonderived carry seeds. Those are independent caller values and remain valid here. @@ -666,7 +648,10 @@ def store( if tracked: public = name.removesuffix("_reported") simulation._drop_values_that_may_depend_on(public, period) - _purge_derived_view(simulation, public) + # Component bases and downstream results also depend on the + # reported source. Core's dependency drops run only on branches, + # so ordinary simulations need the same derived-only cleanup. + _purge_derived_arrays(simulation) result = _PHYSICAL["set"]( holder, period, From ffb62a5a4acfb721a4cef80303655f140e6e5d51 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 16:44:50 -0700 Subject: [PATCH 42/76] Check benefit readers against surviving source inputs Model Core 3.33.12 removal of overlapping raw arrays, and verify reported and public reader values with source selection after each step on every retained copy. Exercise a calculated base and future cached reads. Keep all existing value checks and all 50 differential routes; add zero and distinct-value public cache to compound-setter sequences against main. --- .../core/_social_security_state_backend.py | 93 +++++++++-- ...est_social_security_input_state_machine.py | 152 +++++++++++++++++- .../test_social_security_current_law.py | 49 ++++++ 3 files changed, 280 insertions(+), 14 deletions(-) diff --git a/policyengine_us/tests/core/_social_security_state_backend.py b/policyengine_us/tests/core/_social_security_state_backend.py index e1c8d6a2a24..a332676852f 100644 --- a/policyengine_us/tests/core/_social_security_state_backend.py +++ b/policyengine_us/tests/core/_social_security_state_backend.py @@ -4,7 +4,10 @@ implement every write, builder, copy, replay, extraction and dump route below; the US module's complete function bodies implement the hooks. Only entities, arrays, periods, storage and dependency bookkeeping are synthetic. This tests -input semantics, not benefit calculations or full-model performance. +input semantics and Core's source selection, not country benefit formulas or +full-model performance. A reported -> base -> public formula chain exercises +cache invalidation and scored-year factors. The synthetic uprating index +doubles each year so source changes have exact reader values. """ from __future__ import annotations @@ -55,6 +58,9 @@ class Period: size: int = 1 unit: str = "year" + def __iter__(self): + return iter((self.unit, (self.start.year, self.start.month, 1), self.size)) + @property def stop(self): return Instant(self.start.year + self.size - 1, 12) @@ -264,6 +270,9 @@ def record_unknown_sources(self, *args): def record_derived(self, *args): pass + def prune(self, *args): + pass + class Backend: """Install the unchanged hooks around unchanged Core routes in isolation.""" @@ -288,12 +297,15 @@ def next_sequence(): YEAR="year", MONTH="month", ETERNITY="eternity", + DAY="day", key_period_size=lambda p: p.size, + unit_weight=lambda unit: {"day": 1, "month": 2, "year": 3}[unit], ) commons = types.SimpleNamespace( empty_clone=lambda value: object.__new__(type(value)) ) calculation_frames = ContextVar("state_machine_frames", default=()) + formula_simulation = ContextVar("state_machine_formula", default=None) branch_clone = ContextVar("state_machine_branch_clone", default=None) class BranchClone: @@ -342,6 +354,9 @@ def __init__(self, population, name): def _storage_period(self, period): return as_period(period) + def default_array(self): + return np.zeros(1) + def _raise_if_input_contains_nan(self, array): assert np.isfinite(array).all() @@ -414,6 +429,7 @@ class System: group_entities = [household] cache_blacklist = [] data_modified = False + auto_carry_over_input_variables = True def __init__(self): self.variables = { @@ -425,11 +441,18 @@ def __init__(self): value_type=float, set_input=None, end=None, + requires_computation_after=None, + defined_for=None, + uprating="fixture_index", + calculate_output=None, + quantity_type="flow", + check_uprating_value_type=lambda: None, is_input_variable=lambda name=name: name.endswith("_reported"), ) for name in ( *COMPONENTS, *(f"{name}_reported" for name in COMPONENTS), + *(f"{name}_base" for name in COMPONENTS), ) } @@ -447,6 +470,9 @@ def entities_plural(self): def clone(self): return self + def parameters(self, *args): + return types.SimpleNamespace(gov=types.SimpleNamespace(abolitions={})) + system = System() class Simulation: @@ -459,7 +485,9 @@ def __init__(self, tax_benefit_system=None, populations=None, **kwargs): self._user_input_keys = set() self._user_input_contexts = [] self._inputs_set = 0 + self._input_epoch = 0 self._calculations_started = 0 + self._calculations_in_flight = 0 self.input_variables = [] self.parent_branch = None self.branch_name = "default" @@ -504,11 +532,7 @@ def _get_requested_variables(self): return set() def _drop_computed(self, since=None): - return sum( - holder._drop_computed(since) - for population in self.populations.values() - for holder in population._holders.values() - ) + raise AssertionError("Installed Core cache cleanup was not loaded") def subsample(self, *args): raise NotImplementedError( @@ -516,12 +540,38 @@ def subsample(self, *args): ) def calculate(self, name, period): - return self.get_holder(name).get_array( - as_period(period), self.branch_name - ) + # The country-free fixture has no tracer or concurrent + # formulas. Retain Core's public reader cache behavior while + # its installed _calculate body supplies the actual value. + period = as_period(period) + key = (name, period) + if key not in self._fast_cache: + self._fast_cache[key] = self._calculate(name, period) + return self._fast_cache[key] def _calculate(self, name, period): - return self.calculate(name, period) + raise AssertionError("Installed Core calculation body was not loaded") + + def check_macro_cache(self, *args): + return False + + def _check_period_consistency(self, period, variable): + if period.unit != "year" or period.size != 1: + raise ValueError("A compound reader requires an exact source") + + def _check_for_cycle(self, *args): + pass + + def _run_formula(self, variable, population, period): + if variable.name in COMPONENTS: + return self.calculate(f"{variable.name}_base", period) * 2 + if variable.name.endswith("_base"): + component = variable.name.removesuffix("_base") + return self.calculate(f"{component}_reported", period) + return None + + def _cast_formula_result(self, array, variable): + return np.asarray(array, dtype=float) class DelegatingSimulation(Simulation): def set_input(self, name, period, value): @@ -603,13 +653,21 @@ def load(self): commons=commons, next_sequence_number=next_sequence, _calculation_frames=calculation_frames, + _formula_simulation=formula_simulation, _branch_clone=branch_clone, _BranchClone=BranchClone, _note_unobserved_activity=lambda: None, + _taint_waiting_frames=lambda *args: None, directory_containing=lambda path: None, PreservedUserInput=PreservedInput, NpEncoder=Encoder, ETERNITY="eternity", + YEAR="year", + MONTH="month", + QuantityType=types.SimpleNamespace(STOCK="stock"), + SpiralError=type("SpiralError", (Exception,), {}), + FullTracer=type("FullTracer", (), {}), + get_parameter=lambda *args: lambda instant: 2 ** (instant.year - 2023), Dataset=Dataset, Simulation=Simulation, SimulationBuilder=Builder, @@ -654,6 +712,15 @@ def load(self): if isinstance(node, ast.FunctionDef) ) core_helpers = execute(helper_source, helper_names, namespace) + # Core 3.33.12 cleans overlapping raw seeds after a direct set_input; + # 3.33.2 does this only in dispatch/divide helpers. Keep the reference + # model tied to the installed behavior rather than a version string. + self.cleans_raw_overlaps = "_drop_calculated_over_own_period" in core_helpers + execute( + core / "simulations/simulation.py", + ("_end_order", "_latest_input_key", "_uprating_index_value"), + namespace, + ) sim_methods = ( "set_input", "build_from_dataset", @@ -666,6 +733,12 @@ def load(self): "_get_visible_branch_names", "_get_set_input_periods", "_get_recorded_input_array", + "_calculate", + "_calculation_start", + "_cache_result", + "_may_keep", + "_input_set_meanwhile", + "_drop_computed", ) for name, method in execute( core / "simulations/simulation.py", sim_methods, namespace, "Simulation" diff --git a/policyengine_us/tests/core/test_social_security_input_state_machine.py b/policyengine_us/tests/core/test_social_security_input_state_machine.py index 56a9e40bdf4..db11e51230d 100644 --- a/policyengine_us/tests/core/test_social_security_input_state_machine.py +++ b/policyengine_us/tests/core/test_social_security_input_state_machine.py @@ -4,6 +4,9 @@ period; otherwise its read-time view follows the latest public input. Exact compound totals keep their own period semantics. Copies preserve those sources and public final overrides, while later writes remain isolated. +After every step, actual Core readers must return the independently modeled +reported amount and public formula amount, including the initial scored year +and a later year whose intermediate base may already have been cached. The backend executes installed Core source bodies and the production US hooks without importing the country. Run this file directly for a country-free test: @@ -97,6 +100,76 @@ def reported(self, component, period): return None, False return self.visible(self.public, component, period), True + def reader_source(self, component, period): + """The source an annual reader selects, independently of Core storage.""" + requested = as_period(period) + explicit = { + known: value + for other_component, known in self.keys() + if other_component == component + and as_period(known).start <= requested.start + and (value := self.visible(self.explicit, component, known)) is not None + } + mapping = self.explicit if explicit else self.public + candidates = explicit or { + known: value + for other_component, known in self.keys() + if other_component == component + and as_period(known).start <= requested.start + and (value := self.visible(mapping, component, known)) is not None + } + if period in candidates: + return candidates[period], period, not bool(explicit) + if not candidates: + return None, None, False + # Core first uprates an input starting strictly before the request. + # Only when none exists does carry-over consider a compound input + # starting in the requested year. + earlier = { + known: value + for known, value in candidates.items() + if as_period(known).start < requested.start + } + known = max( + earlier or candidates, + key=lambda text: ( + as_period(text).start, + as_period(text).stop, + text, + ), + ) + return candidates[known], known, not bool(explicit) + + def reader_amount(self, component, period): + value, known, _ = self.reader_source(component, period) + if value is None: + return 0 + return value.amount * 2 ** ( + as_period(period).start.year - as_period(known).start.year + ) + + def drop_overlapping_raw(self, component, reported, period): + """Model Core's newer cleanup of non-input overlapping cache seeds. + + This follows the same rule as main: a raw cache is a carry seed while + it survives, but an overlapping registered set_input may remove it. + Reserved initial sources are registered inputs and are retained. + """ + requested = as_period(period) + for name in ("explicit",) if reported else ("public", "overrides"): + mapping = getattr(self, name) + for key, value in tuple(mapping.items()): + branch, other_component, known = key + candidate = as_period(known) + if ( + branch in self.visible_branches + and other_component == component + and not value.registered + and candidate.start <= requested.stop + and candidate.stop >= requested.start + ): + del mapping[key] + def write( self, component, @@ -107,7 +180,10 @@ def write( route, initial=False, branch_name=None, + cleanup=False, ): + if cleanup: + self.drop_overlapping_raw(component, reported, period) branch_name = self.branch_name if branch_name is None else branch_name key = (branch_name, component, period) value = Input(amount, registered, route) @@ -217,7 +293,15 @@ def set_input(self, slot, component, reported, period, amount, route, data=None) else type(simulation).set_input ) setter(simulation, name, period, [amount]) - model.write(component, reported, period, amount, True, route) + model.write( + component, + reported, + period, + amount, + True, + route, + cleanup=self.backend.cleans_raw_overlaps, + ) @rule( slot=TARGET, @@ -303,8 +387,15 @@ def cache_result( return simulation, model = self.state(slot) component, period = self.write_key(model, component, period, data) - simulation.get_holder(self.name(component, reported)).put_in_cache( - [amount], as_period(period), simulation.branch_name, derived=True + # Cache an actual reader result. An arbitrary fabricated formula + # result could disagree with its own source by construction; this + # rule instead checks that real derived results never become sources. + name = self.name(component, reported) + if as_period(period).size > 1: + period = str(as_period(period).start.year) + value = simulation.calculate(name, as_period(period)) + simulation.get_holder(name).put_in_cache( + value, as_period(period), simulation.branch_name, derived=True ) @rule( @@ -373,6 +464,7 @@ def builder_load( "yaml" if yaml else "builder", initial=True, branch_name="default" if yaml else model.branch_name, + cleanup=self.backend.cleans_raw_overlaps, ) if fresh: self.append(result, model) @@ -426,6 +518,7 @@ def setter_stores_then_raises( True, "failed initial setter" if initial else "failed setter", initial=initial, + cleanup=self.backend.cleans_raw_overlaps, ) @rule(slot=TARGET, branch=st.booleans()) @@ -523,7 +616,7 @@ def reported_values_and_provenance_follow_intent(self): simulation, reported, as_period(period) ) continue - assert actual is not None, (key, expected.route, "missing twin") + assert actual is not None, (key, expected.route, "missing source view") assert float(actual[0]) == expected.amount, ( key, expected.route, @@ -556,6 +649,57 @@ def reported_values_and_provenance_follow_intent(self): float(holder.get_array(period, simulation.branch_name)[0]) == expected.amount ) + # The source view alone cannot expose stale calculated + # intermediates. Exercise Core's actual calculation body on + # the scored years and a later year after every operation. + for year in ("2023", "2024", "2025", "2026"): + expected_amount = model.reader_amount(component, year) + actual_reported = simulation.calculate( + f"{component}_reported", as_period(year) + ) + assert float(actual_reported[0]) == expected_amount, ( + component, + year, + "reported reader", + float(actual_reported[0]), + expected_amount, + ) + source, source_period, expected_automatic = model.reader_source( + component, year + ) + if source is not None: + value, automatic = self.backend.hooks["_effective_value"]( + simulation, + f"{component}_reported", + as_period(source_period), + ) + assert ( + value is not None and float(value.array[0]) == source.amount + ), ( + component, + year, + "selected source", + source_period, + ) + assert automatic == expected_automatic, ( + component, + year, + "selected source provenance", + automatic, + expected_automatic, + ) + override = model.visible(model.overrides, component, year) + expected_public = ( + override.amount if override is not None else expected_amount * 2 + ) + actual_public = simulation.calculate(component, as_period(year)) + assert float(actual_public[0]) == expected_public, ( + component, + year, + "public reader through base and factor", + float(actual_public[0]), + expected_public, + ) for component, period in model.keys(): expected = model.visible(model.overrides, component, period) if expected is None: diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 000f4470656..2f24471fb66 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -463,6 +463,50 @@ def set_input(self, variable, input_period, value): raise AssertionError("The delegating setter did not raise") result[route] = _snapshot(simulation, (2024, 2026), INGRESS_OUTCOMES) del simulation + if route == "setter_failed_after_loading": + # The machine found a raw annual public seed followed by a + # compound public setter that stores and raises. Keep the exact + # shrunk zero sequence and a distinct-value variant in the same + # route. Both checkouts use the installed Core version's source + # cleanup; no reported-name input is substituted on main. + for label, raw_amount, compound_amount in ( + ("zero", 0, 0), + ("nonzero", 12_000, 24_000), + ): + simulation = StoreThenRaiseSimulation( + situation=inputs, tax_benefit_system=system + ) + simulation.get_holder("social_security_retirement").put_in_cache( + np.array([raw_amount]), period(2023), derived=False + ) + before = _snapshot(simulation, (2023, 2024, 2026)) + simulation.store_error = RuntimeError("compound setter stored input") + try: + simulation.set_input( + "social_security_retirement", + "year:2023:2", + [compound_amount], + ) + except RuntimeError as error: + assert error is simulation.store_error + else: + raise AssertionError("The compound setter did not raise") + after = _snapshot(simulation, (2023, 2024, 2026)) + after.update( + _snapshot( + simulation, + ("year:2023:2",), + ("social_security_retirement",), + ) + ) + for stage, snapshot in (("before", before), ("after", after)): + result[route].update( + { + f"raw_cache_compound_{label}_{stage}:{request}": values + for request, values in snapshot.items() + } + ) + del simulation gc.collect() # Subsampling rebuilds holder storage from a dataframe. Its new stores @@ -818,6 +862,11 @@ def differential(tmp_path_factory): f"{completed.stdout[-4000:]}\n{completed.stderr[-8000:]}" ) results[label] = json.loads(destination.read_text()) + print( + f"Social Security differential ({label}): " + f"{len(results[label])} routes, " + f"{sum(len(values) for values in results[label].values())} arrays" + ) return results["current"], results["main"] From 5cff9088f1549af9c13d51da4207590c560b0396 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 17:06:02 -0700 Subject: [PATCH 43/76] Load Core calendar metadata for reader tests The installed input-period sorter uses its module calendar constant. Supply the same literal from the installed source AST so actual reader and uprating checks can execute under both supported Core versions. --- .../tests/core/_social_security_state_backend.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/policyengine_us/tests/core/_social_security_state_backend.py b/policyengine_us/tests/core/_social_security_state_backend.py index a332676852f..4ff5caed4c5 100644 --- a/policyengine_us/tests/core/_social_security_state_backend.py +++ b/policyengine_us/tests/core/_social_security_state_backend.py @@ -716,8 +716,22 @@ def load(self): # 3.33.2 does this only in dispatch/divide helpers. Keep the reference # model tied to the installed behavior rather than a version string. self.cleans_raw_overlaps = "_drop_calculated_over_own_period" in core_helpers + simulation_source = core / "simulations/simulation.py" + # _latest_input_key calls _end_order, which uses Core's calendar + # table even for yearly sources. Load that literal from the same + # installed module rather than maintaining a second calendar table. + day_offsets = next( + node.value + for node in source_tree(simulation_source).body + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "_DAYS_BEFORE_MONTH" + for target in node.targets + ) + ) + namespace["_DAYS_BEFORE_MONTH"] = ast.literal_eval(day_offsets) execute( - core / "simulations/simulation.py", + simulation_source, ("_end_order", "_latest_input_key", "_uprating_index_value"), namespace, ) From 25fde8ace969b30c0f54fe29dc695cb0c9d2f098 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 17:23:30 -0700 Subject: [PATCH 44/76] Report all values on mismatching benefit routes Retain the existing exact equality checks and print every mismatched request before they fail, distinguishing component-reader differences from cached totals and tax results on the main reference. --- .../test_social_security_current_law.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 2f24471fb66..74e1ad939e6 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -927,6 +927,21 @@ def differential(tmp_path_factory): ) def test_legacy_public_ingress_matches_unmodified_main(differential, route): current, main = differential + for request in main[route]: + if request in current[route] and not np.array_equal( + current[route][request], main[route][request] + ): + print( + "Social Security differential mismatch: " + + json.dumps( + { + "route": route, + "request": request, + "current": current[route][request], + "main": main[route][request], + } + ) + ) assert current[route].keys() == main[route].keys() for request in main[route]: np.testing.assert_array_equal( From f0cda5e17bd31f813217c3085d58e735118637a2 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 17:39:17 -0700 Subject: [PATCH 45/76] Verify benefit state machine under both Core versions Run the country-free machine with Core 3.33.2 and 3.33.12, report actual Hypothesis examples and runtime, and enforce the five-minute passing-run budget. Each job runs one foreground file with four-thread caps and a fifteen-minute process cap; existing reference fetches and process guards remain intact. --- .github/workflows/pr.yaml | 97 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 617d18bc6b5..ed01c27ffb8 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -296,3 +296,100 @@ jobs: name: Python tests needs: ReleaseLock uses: ./.github/workflows/python-tests.yaml + + Social-Security-Input-Machine: + needs: ReleaseLock + name: Social Security input machine (Core ${{ matrix.core }}) + runs-on: ubuntu-latest + timeout-minutes: 20 + strategy: + fail-fast: false + matrix: + core: ["3.33.2", "3.33.12"] + env: + PYTHONUNBUFFERED: "1" + OMP_NUM_THREADS: "4" + OPENBLAS_NUM_THREADS: "4" + MKL_NUM_THREADS: "4" + VECLIB_MAXIMUM_THREADS: "4" + NUMEXPR_NUM_THREADS: "4" + OMP_WAIT_POLICY: PASSIVE + MACHINE_CORE_VERSION: ${{ matrix.core }} + steps: + - name: Checkout repo + uses: actions/checkout@v6 + - name: Update system certificate authorities + uses: ./.github/actions/setup-system-ca + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.13" + - name: Install uv + uses: astral-sh/setup-uv@v8.1.0 + with: + version: "0.13.4" + - name: Install country-free machine dependencies + run: | + uv venv --python "$(command -v python)" .machine-venv + uv pip install --python .machine-venv/bin/python "policyengine-core==$MACHINE_CORE_VERSION" "hypothesis==6.168.5" + git diff --exit-code -- pyproject.toml uv.lock + - name: Run one state-machine file with example statistics + run: | + timeout 900 .machine-venv/bin/python - <<'PYTHON' + import importlib.util + import json + import os + import sys + import time + import unittest + from importlib.metadata import version + from pathlib import Path + + from hypothesis.statistics import collector, describe_statistics + + core_version = version("policyengine-core") + assert core_version == os.environ["MACHINE_CORE_VERSION"], core_version + print(f"Core version: {core_version}", flush=True) + path = Path("policyengine_us/tests/core/test_social_security_input_state_machine.py") + spec = importlib.util.spec_from_file_location("social_security_ci_machine", path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + statistics = [] + + def record(stats): + statistics.append(stats) + print(describe_statistics(stats), flush=True) + + started = time.monotonic() + with collector.with_value(record): + result = unittest.TextTestRunner(verbosity=2).run( + unittest.defaultTestLoader.loadTestsFromModule(module) + ) + elapsed = time.monotonic() - started + summary = { + "core_version": core_version, + "elapsed_seconds": elapsed, + "tests_run": result.testsRun, + "failures": len(result.failures), + "errors": len(result.errors), + "configured_examples": module.TestSocialSecurityInputStateMachine.settings.max_examples, + "configured_steps": module.TestSocialSecurityInputStateMachine.settings.stateful_step_count, + "statistics": statistics, + } + Path(f"social-security-machine-{core_version}.json").write_text( + json.dumps(summary, indent=2, default=str) + "\n" + ) + print(json.dumps({key: value for key, value in summary.items() if key != "statistics"}), flush=True) + if not result.wasSuccessful(): + raise SystemExit(1) + if elapsed >= 300: + raise SystemExit("State-machine file exceeded its five-minute budget") + PYTHON + - name: Upload machine runtime and example statistics + if: always() + uses: actions/upload-artifact@v4 + with: + name: social-security-machine-core-${{ matrix.core }} + path: social-security-machine-${{ matrix.core }}.json + if-no-files-found: error From 70b4949b9e0b0c1d9f9d0c3a6567bfacb87d1753 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 21:42:38 -0700 Subject: [PATCH 46/76] Remove unapproved Social Security CI job --- .github/workflows/pr.yaml | 97 --------------------------------------- 1 file changed, 97 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index ed01c27ffb8..617d18bc6b5 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -296,100 +296,3 @@ jobs: name: Python tests needs: ReleaseLock uses: ./.github/workflows/python-tests.yaml - - Social-Security-Input-Machine: - needs: ReleaseLock - name: Social Security input machine (Core ${{ matrix.core }}) - runs-on: ubuntu-latest - timeout-minutes: 20 - strategy: - fail-fast: false - matrix: - core: ["3.33.2", "3.33.12"] - env: - PYTHONUNBUFFERED: "1" - OMP_NUM_THREADS: "4" - OPENBLAS_NUM_THREADS: "4" - MKL_NUM_THREADS: "4" - VECLIB_MAXIMUM_THREADS: "4" - NUMEXPR_NUM_THREADS: "4" - OMP_WAIT_POLICY: PASSIVE - MACHINE_CORE_VERSION: ${{ matrix.core }} - steps: - - name: Checkout repo - uses: actions/checkout@v6 - - name: Update system certificate authorities - uses: ./.github/actions/setup-system-ca - - name: Set up Python - uses: actions/setup-python@v6 - with: - python-version: "3.13" - - name: Install uv - uses: astral-sh/setup-uv@v8.1.0 - with: - version: "0.13.4" - - name: Install country-free machine dependencies - run: | - uv venv --python "$(command -v python)" .machine-venv - uv pip install --python .machine-venv/bin/python "policyengine-core==$MACHINE_CORE_VERSION" "hypothesis==6.168.5" - git diff --exit-code -- pyproject.toml uv.lock - - name: Run one state-machine file with example statistics - run: | - timeout 900 .machine-venv/bin/python - <<'PYTHON' - import importlib.util - import json - import os - import sys - import time - import unittest - from importlib.metadata import version - from pathlib import Path - - from hypothesis.statistics import collector, describe_statistics - - core_version = version("policyengine-core") - assert core_version == os.environ["MACHINE_CORE_VERSION"], core_version - print(f"Core version: {core_version}", flush=True) - path = Path("policyengine_us/tests/core/test_social_security_input_state_machine.py") - spec = importlib.util.spec_from_file_location("social_security_ci_machine", path) - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - statistics = [] - - def record(stats): - statistics.append(stats) - print(describe_statistics(stats), flush=True) - - started = time.monotonic() - with collector.with_value(record): - result = unittest.TextTestRunner(verbosity=2).run( - unittest.defaultTestLoader.loadTestsFromModule(module) - ) - elapsed = time.monotonic() - started - summary = { - "core_version": core_version, - "elapsed_seconds": elapsed, - "tests_run": result.testsRun, - "failures": len(result.failures), - "errors": len(result.errors), - "configured_examples": module.TestSocialSecurityInputStateMachine.settings.max_examples, - "configured_steps": module.TestSocialSecurityInputStateMachine.settings.stateful_step_count, - "statistics": statistics, - } - Path(f"social-security-machine-{core_version}.json").write_text( - json.dumps(summary, indent=2, default=str) + "\n" - ) - print(json.dumps({key: value for key, value in summary.items() if key != "statistics"}), flush=True) - if not result.wasSuccessful(): - raise SystemExit(1) - if elapsed >= 300: - raise SystemExit("State-machine file exceeded its five-minute budget") - PYTHON - - name: Upload machine runtime and example statistics - if: always() - uses: actions/upload-artifact@v4 - with: - name: social-security-machine-core-${{ matrix.core }} - path: social-security-machine-${{ matrix.core }}.json - if-no-files-found: error From 2ecebb15c1fa0c0847377660a89c8eca6e5baad1 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 21:50:53 -0700 Subject: [PATCH 47/76] Compare fresh benefit formulas and assert intended cache differences --- .../test_social_security_current_law.py | 165 +++++++++++++++++- 1 file changed, 162 insertions(+), 3 deletions(-) diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 74e1ad939e6..7e3c3b53ac1 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -2,6 +2,9 @@ The same probe runs in separate Python processes against this checkout and main. No variable definitions or input names are substituted in the reference. +After benefit sources change, both probes discard only formula-derived caches +to compare freshly calculated current law. Separate intended-difference tests +check the live head's values before that cleanup at main's eight stale routes. An unmodified pre-feature main checkout is required. Set PE_US_MAIN_CHECKOUT to that checkout, or make the reviewed main commit or a pre-feature main ref available locally (the fixture archives it without changing either checkout). @@ -53,6 +56,12 @@ ) AMOUNTS = ((12_000, 0, 1_800, 600), (0, 6_000, 1_200, 300)) REVIEWED_MAIN_SHA = "8986e05cab65d56435a23df1910e2d35f2c5f5bf" +# Hand-computed per-person growth from calibration.gov.cbo.social_security +# divided by calibration.gov.census.populations.total in each year. The IRS +# SOI series follows this CBO growth; Core multiplies float32 benefit arrays. +UPRATING_2024_TO_2026 = np.float32( + (1_665_942_000_000 / 352_304_000) / (1_454_424_000_000 / 346_588_000) +) def situation(*, include_benefits=True): @@ -108,8 +117,34 @@ def _snapshot(simulation, periods, variables=BENEFITS): } +def _clear_derived_caches(simulation): + """Compare formula results for the same surviving supplied/raw sources. + + Runtime invalidation is checked before any cleanup by the live-value + assertions, state machine and input regressions. This differential + separately compares the country formulas after source writes, including + main's ordinary simulations. + """ + for population in simulation.populations.values(): + for holder in population._holders.values(): + memory = holder._memory_storage + for branch, period in tuple(memory.get_known_branch_periods()): + if memory.is_derived(period, branch): + memory._arrays.pop(f"{branch}:{period}", None) + memory._unmark_dropped_keys() + memory._stop_sharing_dropped_keys() + memory._forget_dropped_numbers() + disk = holder._disk_storage + if disk is not None: + for branch, period in tuple(disk.get_known_branch_periods()): + if disk.is_derived(period, branch): + disk._files.pop(f"{branch}_{period}", None) + disk._forget_dropped_keys() + simulation._fast_cache = {} + + def _write_h5(path): - """Variable-centric v2pop-style H5, with public columns in its scored year.""" + """Public-column H5, with source values in its scored year.""" import h5py ids = np.array([1, 2]) @@ -188,6 +223,15 @@ def set_input(self, variable, input_period, value): else None ) result = {} + intended_differences = {} + + def record_live_value(route, simulation, variable, year): + # Capture before normalization so the intended-difference assertions + # independently require the live head to invalidate stale values. + intended_differences[route] = np.asarray( + simulation.calculate(variable, year) + ).tolist() + for label, wrapper in ( ("simulation", Simulation), ("microsimulation", HouseholdMicrosimulation), @@ -227,6 +271,10 @@ def set_input(self, variable, input_period, value): simulation, (2024, 2025, 2026, 2027) ) simulation.set_input("social_security_retirement", 2026, [36_000, 0]) + record_live_value( + "complete_export_rewrite", simulation, "social_security", 2026 + ) + _clear_derived_caches(simulation) result["complete_export_rewrite"] = _snapshot(simulation, (2026, 2027)) else: payload = export_source_inputs(simulation) @@ -242,6 +290,10 @@ def set_input(self, variable, input_period, value): simulation, (2024, 2025, 2026, 2027) ) reloaded.set_input("social_security_retirement", 2026, [36_000, 0]) + record_live_value( + "complete_export_rewrite", reloaded, "social_security", 2026 + ) + _clear_derived_caches(reloaded) result["complete_export_rewrite"] = _snapshot(reloaded, (2026, 2027)) del reloaded del simulation @@ -291,6 +343,10 @@ def set_input(self, variable, input_period, value): } result[f"extraction_future_reload_{label}"] = extracted_future reloaded.set_input("social_security_retirement", 2024, [36_000]) + record_live_value( + f"extraction_rewrite_{label}", reloaded, "social_security", 2024 + ) + _clear_derived_caches(reloaded) result[f"extraction_rewrite_{label}"] = _snapshot( reloaded, (2024, 2026, 2027), (*BENEFITS, *DOWNSTREAM) ) @@ -418,6 +474,13 @@ def set_input(self, variable, input_period, value): reloaded, (2024, 2026), INGRESS_OUTCOMES ) reloaded.set_input("social_security_retirement", 2024, [36_000]) + record_live_value( + "dataset_builder_extraction_rewrite", + reloaded, + "social_security_retirement", + 2026, + ) + _clear_derived_caches(reloaded) result["dataset_builder_extraction_rewrite"] = _snapshot( reloaded, (2024, 2026), INGRESS_OUTCOMES ) @@ -491,6 +554,9 @@ def set_input(self, variable, input_period, value): assert error is simulation.store_error else: raise AssertionError("The compound setter did not raise") + if label == "nonzero": + record_live_value(route, simulation, "social_security", 2023) + _clear_derived_caches(simulation) after = _snapshot(simulation, (2023, 2024, 2026)) after.update( _snapshot( @@ -562,6 +628,13 @@ def set_input(self, variable, input_period, value): ) assert result["dump_restore_legacy_public"] == live restored.set_input("social_security_retirement", 2024, [24_000]) + record_live_value( + "dump_restore_legacy_public_override", + restored, + "social_security_retirement", + 2026, + ) + _clear_derived_caches(restored) result["dump_restore_legacy_public_override"] = _snapshot( restored, (2024, 2026), INGRESS_OUTCOMES ) @@ -575,6 +648,14 @@ def set_input(self, variable, input_period, value): assert result["dump_restore_both_names"] == live for candidate in (simulation, restored): candidate.set_input("social_security_retirement", 2024, [24_000]) + if candidate is restored: + record_live_value( + "dump_restore_both_names_override", + restored, + "social_security_retirement", + 2026, + ) + _clear_derived_caches(candidate) result["dump_restore_both_names_override"] = _snapshot( restored, (2024, 2026), INGRESS_OUTCOMES ) @@ -663,6 +744,9 @@ def set_input(self, variable, input_period, value): else simulation.get_branch("public_override") ) copied.set_input("social_security_retirement", 2026, [24_000, 0]) + if route == "clone": + record_live_value(route, copied, "social_security", 2026) + _clear_derived_caches(copied) result[route] = _snapshot(copied, (2026, 2027)) assert _snapshot(simulation, (2024, 2026)) == parent assert simulation.input_variables == original_inputs @@ -743,6 +827,9 @@ def set_input(self, variable, input_period, value): rtol=1e-7, atol=0, ) + intended_destination = os.environ.get("PE_US_DIFFERENTIAL_LIVE_VALUES") + if intended_destination: + Path(intended_destination).write_text(json.dumps(intended_differences)) return result @@ -826,14 +913,17 @@ def differential(tmp_path_factory): current = Path(__file__).resolve().parents[3] reference = _main_checkout(current, directory) results = {} + live_values = {} dump_directory = directory / "dumps" dump_directory.mkdir() # The main probe publishes a real legacy dump for the current probe. for label, checkout in (("main", reference), ("current", current)): destination = directory / f"{label}.json" + live_destination = directory / f"{label}-live.json" environment = dict(os.environ) environment["PYTHONPATH"] = str(checkout) environment["PE_US_DIFFERENTIAL_DUMPS"] = str(dump_directory) + environment["PE_US_DIFFERENTIAL_LIVE_VALUES"] = str(live_destination) for variable in ( "OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", @@ -862,12 +952,13 @@ def differential(tmp_path_factory): f"{completed.stdout[-4000:]}\n{completed.stderr[-8000:]}" ) results[label] = json.loads(destination.read_text()) + live_values[label] = json.loads(live_destination.read_text()) print( f"Social Security differential ({label}): " f"{len(results[label])} routes, " f"{sum(len(values) for values in results[label].values())} arrays" ) - return results["current"], results["main"] + return results["current"], results["main"], live_values["current"] @pytest.mark.parametrize( @@ -926,7 +1017,7 @@ def differential(tmp_path_factory): ), ) def test_legacy_public_ingress_matches_unmodified_main(differential, route): - current, main = differential + current, main, _ = differential for request in main[route]: if request in current[route] and not np.array_equal( current[route][request], main[route][request] @@ -951,6 +1042,74 @@ def test_legacy_public_ingress_matches_unmodified_main(differential, route): ) +@pytest.mark.parametrize( + "route, expected", + ( + # Intended difference from main: its cached total remains 24,000. + ("extraction_rewrite_simulation", [36_000 + 0 + 0 + 0]), + # Intended difference from main: its cached total remains 24,000. + ("extraction_rewrite_microsimulation", [36_000 + 0 + 0 + 0]), + # Intended difference from main: its first total remains + # 26,704.431640625 after the 24,000 override is replaced by 36,000. + ( + "complete_export_rewrite", + [ + np.float32(36_000) + + np.float32(1_800) * UPRATING_2024_TO_2026 + + np.float32(600) * UPRATING_2024_TO_2026, + np.float32(6_000) * UPRATING_2024_TO_2026 + + np.float32(1_200) * UPRATING_2024_TO_2026 + + np.float32(300) * UPRATING_2024_TO_2026, + ], + ), + # Intended difference from main: its cached 2026 retirement remains + # 27,044.31640625, uprated from the previous 24,000 source. + ( + "dataset_builder_extraction_rewrite", + [np.float32(36_000) * UPRATING_2024_TO_2026], + ), + # Intended difference from main: its cached 2023 total remains + # 12,000 after the nonzero raw seed's compound setter stores 24,000. + ("setter_failed_after_loading", [24_000 + 0 + 0 + 0]), + # Intended difference from main: its cached 2026 retirement remains + # 13,522.158203125, uprated from the restored 12,000 source. + ( + "dump_restore_legacy_public_override", + [np.float32(24_000) * UPRATING_2024_TO_2026], + ), + # Intended difference from main: its cached 2026 retirement remains + # 13,522.158203125, uprated from the restored 12,000 source. + ( + "dump_restore_both_names_override", + [np.float32(24_000) * UPRATING_2024_TO_2026], + ), + # Intended difference from main: its cloned 2026 first total remains + # 16,226.58984375 after retirement is overridden to 24,000. + ( + "clone", + [ + np.float32(24_000) + + np.float32(1_800) * UPRATING_2024_TO_2026 + + np.float32(600) * UPRATING_2024_TO_2026, + np.float32(6_000) * UPRATING_2024_TO_2026 + + np.float32(1_200) * UPRATING_2024_TO_2026 + + np.float32(300) * UPRATING_2024_TO_2026, + ], + ), + ), +) +def test_intended_difference_from_main_uses_fresh_live_value( + differential, route, expected +): + """A benefit source written after calculation replaces its stale result.""" + _, _, live_values = differential + np.testing.assert_array_equal( + live_values[route], + expected, + err_msg=f"Intended difference from main did not use the latest source: {route}", + ) + + if __name__ == "__main__": if len(sys.argv) != 3 or sys.argv[1] != "--probe": raise SystemExit("Usage: test_social_security_current_law.py --probe OUTPUT") From d91e28f61c7c4f72a599a60e886484f4916d5fd8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 23:00:05 -0700 Subject: [PATCH 48/76] Exercise inherited sources across branch builder rewrites --- ...est_social_security_input_state_machine.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/policyengine_us/tests/core/test_social_security_input_state_machine.py b/policyengine_us/tests/core/test_social_security_input_state_machine.py index db11e51230d..ac52ea66cf6 100644 --- a/policyengine_us/tests/core/test_social_security_input_state_machine.py +++ b/policyengine_us/tests/core/test_social_security_input_state_machine.py @@ -537,6 +537,37 @@ def copy_simulation(self, slot, branch): ) self.append(copied, copied_model) + @rule( + slot=TARGET, + component=COMPONENT, + period=PERIOD, + amount=AMOUNT, + replacement=AMOUNT, + ) + def branch_initial_rewrite(self, slot, component, period, amount, replacement): + """Compose an inherited final override with a branch's initial rewrite.""" + _, model = self.state(slot) + if model.visible(model.explicit, component, period) is not None: + # A same-period explicit twin intentionally retains the public + # override; this composition exercises initial source routing. + return + # Independent rules rarely choose the same component and period on + # a newly created branch. Keep that composition generated, checking + # the parent and every earlier copy after each intermediate step. + self.set_input(slot, component, False, period, amount, "simulation") + self.reported_values_and_provenance_follow_intent() + self.copy_simulation(slot, branch=True) + self.reported_values_and_provenance_follow_intent() + self.builder_load( + len(self.states) - 1, + component, + False, + period, + replacement, + fresh=False, + yaml=False, + ) + @rule(slot=TARGET) def dump_restore(self, slot): simulation, model = self.state(slot) From d52f0c0b7e78e6ed517dc67993e84bb686b00e6c Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 25 Oct 2026 23:27:58 -0700 Subject: [PATCH 49/76] Bound composed branch coverage per generated example --- .../tests/core/test_social_security_input_state_machine.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/policyengine_us/tests/core/test_social_security_input_state_machine.py b/policyengine_us/tests/core/test_social_security_input_state_machine.py index ac52ea66cf6..768432ee89a 100644 --- a/policyengine_us/tests/core/test_social_security_input_state_machine.py +++ b/policyengine_us/tests/core/test_social_security_input_state_machine.py @@ -24,7 +24,7 @@ from pathlib import Path from hypothesis import HealthCheck, settings, strategies as st -from hypothesis.stateful import RuleBasedStateMachine, invariant, rule +from hypothesis.stateful import RuleBasedStateMachine, invariant, precondition, rule spec = importlib.util.spec_from_file_location( @@ -237,6 +237,7 @@ def __init__(self): self.states = [(self.backend.DelegatingSimulation(), Model())] self.temp = tempfile.TemporaryDirectory(prefix="ss-input-machine-") self.dump_count = 0 + self.branch_initial_rewrite_done = False def teardown(self): self.temp.cleanup() @@ -537,6 +538,7 @@ def copy_simulation(self, slot, branch): ) self.append(copied, copied_model) + @precondition(lambda self: not self.branch_initial_rewrite_done) @rule( slot=TARGET, component=COMPONENT, @@ -554,6 +556,8 @@ def branch_initial_rewrite(self, slot, component, period, amount, replacement): # Independent rules rarely choose the same component and period on # a newly created branch. Keep that composition generated, checking # the parent and every earlier copy after each intermediate step. + # One composed rewrite bounds retained-copy growth; the independent + # primitive rules can still compose before and after it on any state. self.set_input(slot, component, False, period, amount, "simulation") self.reported_values_and_provenance_follow_intent() self.copy_simulation(slot, branch=True) @@ -567,6 +571,7 @@ def branch_initial_rewrite(self, slot, component, period, amount, replacement): fresh=False, yaml=False, ) + self.branch_initial_rewrite_done = True @rule(slot=TARGET) def dump_restore(self, slot): From 582625618d5d88f1fec8aef7d392cc584c1d5698 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 01:20:59 -0700 Subject: [PATCH 50/76] Observe compound source refresh across Core versions --- .../test_social_security_current_law.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 7e3c3b53ac1..c103a29f4e1 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -59,6 +59,9 @@ # Hand-computed per-person growth from calibration.gov.cbo.social_security # divided by calibration.gov.census.populations.total in each year. The IRS # SOI series follows this CBO growth; Core multiplies float32 benefit arrays. +UPRATING_2023_TO_2026 = np.float32( + (1_665_942_000_000 / 352_304_000) / (1_379_268_000_000 / 342_835_000) +) UPRATING_2024_TO_2026 = np.float32( (1_665_942_000_000 / 352_304_000) / (1_454_424_000_000 / 346_588_000) ) @@ -555,7 +558,12 @@ def record_live_value(route, simulation, variable, year): else: raise AssertionError("The compound setter did not raise") if label == "nonzero": - record_live_value(route, simulation, "social_security", 2023) + # Core 3.33.2 retains the exact 2023 raw seed. Both Core + # versions select the replacement compound source in + # 2026; main still retains its cached retirement there. + record_live_value( + route, simulation, "social_security_retirement", 2026 + ) _clear_derived_caches(simulation) after = _snapshot(simulation, (2023, 2024, 2026)) after.update( @@ -1068,9 +1076,12 @@ def test_legacy_public_ingress_matches_unmodified_main(differential, route): "dataset_builder_extraction_rewrite", [np.float32(36_000) * UPRATING_2024_TO_2026], ), - # Intended difference from main: its cached 2023 total remains - # 12,000 after the nonzero raw seed's compound setter stores 24,000. - ("setter_failed_after_loading", [24_000 + 0 + 0 + 0]), + # Intended difference from main: its cached 2026 retirement remains + # 14,104.5751953125, uprated from the previous 12,000 raw source. + ( + "setter_failed_after_loading", + [np.float32(24_000) * UPRATING_2023_TO_2026], + ), # Intended difference from main: its cached 2026 retirement remains # 13,522.158203125, uprated from the restored 12,000 source. ( From 2a289b5347a4bba194bb5058c17305355c3f1973 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 03:37:31 -0700 Subject: [PATCH 51/76] Preserve Social Security helper precedence and bounded refreshes --- changelog.d/ss-benefit-levers.added.md | 2 +- .../core/test_social_security_real_core.py | 349 ++++++++++++++++++ .../tools/social_security_inputs.py | 100 ++++- 3 files changed, 431 insertions(+), 20 deletions(-) create mode 100644 policyengine_us/tests/core/test_social_security_real_core.py diff --git a/changelog.d/ss-benefit-levers.added.md b/changelog.d/ss-benefit-levers.added.md index 31b9fa30ce7..f04f3a13901 100644 --- a/changelog.d/ss-benefit-levers.added.md +++ b/changelog.d/ss-benefit-levers.added.md @@ -1 +1 @@ -Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs remain independent sources, exposed through a read-time reported view without automatic twin arrays or provenance stamps. Initial loading routes sources so same-year benefit factors apply; later public writes remain final overrides. Explicit reported inputs retain priority, while exact compound totals keep their period semantics. Builders, partial loading failures, copies, cache replay, extraction, dump restoration and subsampling preserve successful source inputs. A state machine searches combinations of write and copy routes. +Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs remain independent sources, exposed through a read-time reported view without automatic twin arrays or provenance stamps. Initial loading routes sources so same-year benefit factors apply; later public writes remain final overrides. Explicit reported inputs retain priority, while exact compound totals keep their period semantics. Compound helpers retain initial-input precedence. Source refreshes use Core's dependency sequence cutoff, preserving raw caller values, supplied inputs and earlier formula results. Builders, partial loading failures, copies, cache replay, extraction, dump restoration and subsampling preserve successful source inputs. A state machine searches combinations of write and copy routes. diff --git a/policyengine_us/tests/core/test_social_security_real_core.py b/policyengine_us/tests/core/test_social_security_real_core.py new file mode 100644 index 00000000000..b266e7a0a78 --- /dev/null +++ b/policyengine_us/tests/core/test_social_security_real_core.py @@ -0,0 +1,349 @@ +"""Exercise source hooks on real Core storage, helpers and simulations. + +This file deliberately does not import the country. Its tiny TaxBenefitSystem +uses the production hooks and installed Core implementations without stubs. +Run it separately from country-importing tests. +""" + +from __future__ import annotations + +import importlib +import sys +import types +from pathlib import Path + +import numpy as np +import pytest +from policyengine_core.entities import build_entity +from policyengine_core.holders import Holder, helpers +from policyengine_core.holders.helpers import ( + set_input_dispatch_by_period, + set_input_divide_by_period, +) +from policyengine_core.parameters import ParameterNode +from policyengine_core.periods import period +from policyengine_core.simulations import Simulation, SimulationBuilder +from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from policyengine_core.tools import simulation_dumper +from policyengine_core.variables import Variable + + +COUNTRY_ROOT = Path(__file__).resolve().parents[2] +COMPONENTS = ( + "social_security_retirement", + "social_security_disability", + "social_security_survivors", + "social_security_dependents", +) +PUBLIC = COMPONENTS[0] +hooks = None + + +@pytest.fixture(scope="module", autouse=True) +def installed_hooks(): + """Install only during execution and restore any pre-existing hooks. + + Collection never creates synthetic country packages. This also keeps + grouped CI collection safe if another file imports the actual country. + """ + global hooks + targets = (Holder, Simulation, SimulationBuilder, helpers, simulation_dumper) + originals = {target: dict(vars(target)) for target in targets} + country_modules = { + name: module + for name, module in sys.modules.items() + if name == "policyengine_us" or name.startswith("policyengine_us.") + } + synthetic_country = "policyengine_us" not in sys.modules + try: + if synthetic_country: + # Make the production tools importable without executing the + # country __init__, while retaining normal importlib.reload. + for name, path in ( + ("policyengine_us", COUNTRY_ROOT), + ("policyengine_us.tools", COUNTRY_ROOT / "tools"), + ): + package = types.ModuleType(name) + package.__path__ = [str(path)] + sys.modules[name] = package + hooks = importlib.import_module("policyengine_us.tools.social_security_inputs") + assert hooks.SOCIAL_SECURITY_COMPONENTS == COMPONENTS + hooks.install_social_security_input_compatibility() + yield + finally: + for target, previous in originals.items(): + for name in set(vars(target)) - previous.keys(): + delattr(target, name) + for name, value in previous.items(): + if vars(target).get(name) is not value: + setattr(target, name, value) + if synthetic_country: + for name in tuple(sys.modules): + if name == "policyengine_us" or name.startswith("policyengine_us."): + sys.modules.pop(name) + sys.modules.update(country_modules) + hooks = None + + +def make_system(*, enabled=True, factor=1): + person = build_entity("person", "people", "Person", is_person=True) + household = build_entity( + "household", + "households", + "Household", + roles=[{"key": "member", "plural": "members"}], + ) + system = TaxBenefitSystem([person, household]) + system._social_security_input_compatibility = enabled + system.auto_carry_over_input_variables = True + system.parameters = ParameterNode( + "", + data={ + "fixture_index": { + "values": { + "2023-01-01": 1, + "2024-01-01": 2, + "2025-01-01": 4, + "2026-01-01": 8, + } + } + }, + ) + for component in COMPONENTS: + for suffix in ("_reported", "_base", ""): + name = component + suffix + attributes = { + "value_type": float, + "label": name, + "entity": person, + "definition_period": "year", + "__module__": __name__, + } + if suffix == "_reported": + attributes["uprating"] = "fixture_index" + elif suffix == "_base": + + def formula(person, period, parameters, component=component): + return person(component + "_reported", period) + + attributes["formula"] = formula + elif not enabled: + # Main's public source has uprating and no formula; the + # adapter moves uprating to reported and makes public a formula. + attributes["uprating"] = "fixture_index" + else: + + def formula(person, period, parameters, component=component): + return person(component + "_base", period) * factor + + attributes["formula"] = formula + system.add_variable(type(name, (Variable,), attributes)) + + class independent_input(Variable): + label = "Independent input" + value_type = float + entity = person + definition_period = "year" + + class independent_formula(Variable): + label = "Independent formula" + value_type = float + entity = person + definition_period = "year" + + def formula(person, period, parameters): + return person("independent_input", period) + 7 + + system.add_variables(independent_input, independent_formula) + return system + + +def fresh(inputs=None, *, enabled=True, factor=1): + values = {"independent_input": {"2023": [70.0], "2024": [70.0]}} + values.update(inputs or {}) + return SimulationBuilder().build_from_variables( + make_system(enabled=enabled, factor=factor), values + ) + + +def scalar(simulation, name, requested): + return float(simulation.calculate(name, requested)[0]) + + +def stored(simulation, name, requested): + return simulation.get_holder(name).get_array( + period(requested), simulation.branch_name + ) + + +def earliest_source_dependency(simulation, component=PUBLIC, requested="2024"): + history = simulation._get_store_history() + candidates = [ + history.earliest_dependency(name, period(requested)) + for name in (component, component + "_reported") + ] + return min((number for number in candidates if number is not None), default=None) + + +def sequence_number(simulation, name, requested): + holder = simulation.get_holder(name) + return holder._memory_storage._sequence_numbers[ + f"{simulation.branch_name}:{period(requested)}" + ] + + +HELPER_ROUTES = ( + ("divide", "year:2024:2", ("2024", "2025")), + ("dispatch", "year:2024:2", ("2024", "2025")), + ("ordinary_month", "month:2024-01:24", ("2024", "2025")), + ("ordinary_year", "year:2024:2", ("2024", "2025")), + ( + "rolling_month", + "month:2024-07:24", + ("year:2024-07", "year:2025-07"), + ), +) + + +def apply_helper(simulation, component, route, compound): + value = np.array([36_000.0]) + if route == "divide": + set_input_divide_by_period( + simulation.get_holder(component), period(compound), value + ) + elif route == "dispatch": + set_input_dispatch_by_period( + simulation.get_holder(component), period(compound), value + ) + else: + simulation.set_input(component, compound, value) + + +@pytest.mark.parametrize("component", COMPONENTS) +@pytest.mark.parametrize("route,compound,annual", HELPER_ROUTES) +@pytest.mark.parametrize("initial", (False, True)) +def test_compound_helpers_match_real_core(component, route, compound, annual, initial): + inputs = {component: {annual[0]: [12_000.0]}} if initial else {} + control = fresh(inputs, enabled=False) + simulation = fresh(inputs) + apply_helper(control, component, route, compound) + apply_helper(simulation, component, route, compound) + actual = [scalar(simulation, component, year) for year in annual] + expected = [scalar(control, component, year) for year in annual] + assert actual == expected + # The source view must preserve the same allocation too, even where the + # initial source stays hidden from the public formula's ordinary storage. + assert [scalar(simulation, component + "_reported", year) for year in annual] == ( + expected + ) + + +def test_helper_keeps_initial_source_invisible_to_public_formula(): + simulation = fresh({PUBLIC: {"2024": [12_000.0]}}, factor=2) + apply_helper(simulation, PUBLIC, "dispatch", "year:2024:2") + assert stored(simulation, PUBLIC, "2024") is None + assert scalar(simulation, PUBLIC + "_reported", "2024") == 12_000 + assert scalar(simulation, PUBLIC, "2024") == 24_000 + + +def test_low_level_branch_write_preserves_raw_caller_seed(): + simulation = fresh({PUBLIC: {"2024": [12_000.0]}}) + branch = simulation.get_branch("child") + assert scalar(branch, PUBLIC, "2026") == 48_000 + raw = branch.get_holder("independent_input") + raw.put_in_cache( + np.array([99.0]), period("2025"), branch.branch_name, derived=False + ) + branch.get_holder(PUBLIC)._set( + period("2024"), np.array([24_000.0]), branch.branch_name, is_input=True + ) + np.testing.assert_array_equal(stored(branch, "independent_input", "2025"), [99]) + assert scalar(branch, PUBLIC, "2026") == 96_000 + + +def test_bounded_refresh_preserves_derived_value_before_since(): + simulation = fresh() + assert scalar(simulation, "independent_formula", "2023") == 77 + before = stored(simulation, "independent_formula", "2023") + number = sequence_number(simulation, "independent_formula", "2023") + simulation.set_input(PUBLIC, "2024", [12_000.0]) + assert scalar(simulation, PUBLIC + "_reported", "2026") == 48_000 + assert number < earliest_source_dependency(simulation) + simulation.set_input(PUBLIC, "2024", [24_000.0]) + assert stored(simulation, "independent_formula", "2023") is before + assert scalar(simulation, PUBLIC + "_reported", "2026") == 96_000 + + +def test_bounded_refresh_intentionally_drops_derived_value_after_since(): + simulation = fresh({PUBLIC: {"2024": [12_000.0]}}) + assert scalar(simulation, PUBLIC + "_reported", "2026") == 48_000 + assert scalar(simulation, "independent_formula", "2024") == 77 + assert sequence_number( + simulation, "independent_formula", "2024" + ) >= earliest_source_dependency(simulation) + epoch = simulation._input_epoch + simulation.set_input(PUBLIC, "2024", [24_000.0]) + # Intended over-approximation: RULING-10400-refresh-scope-2026-10-26, + # rule 5 permits unrelated formula values numbered at or after since. + assert stored(simulation, "independent_formula", "2024") is None + assert ("independent_formula", period("2024")) not in simulation._fast_cache + assert simulation._input_epoch > epoch + assert scalar(simulation, "independent_formula", "2024") == 77 + + +def test_bounded_refresh_considers_reported_twin_history(): + simulation = fresh({PUBLIC + "_reported": {"2024": [6_000.0]}}) + assert scalar(simulation, "independent_formula", "2024") == 77 + history = simulation._get_store_history() + assert history.earliest_dependency(PUBLIC, period("2024")) is None + assert history.earliest_dependency(PUBLIC + "_reported", period("2024")) is not None + simulation.set_input(PUBLIC, "2024", [12_000.0]) + assert stored(simulation, "independent_formula", "2024") is None + assert scalar(simulation, PUBLIC + "_reported", "2024") == 6_000 + + +@pytest.mark.parametrize("registered", (False, True), ids=("raw_seed", "input")) +def test_bounded_refresh_preserves_caller_values_after_since(registered): + simulation = fresh({PUBLIC: {"2024": [12_000.0]}}) + scalar(simulation, PUBLIC + "_reported", "2026") + holder = simulation.get_holder("independent_input") + if registered: + simulation.set_input("independent_input", "2025", [99.0]) + else: + holder.put_in_cache(np.array([99.0]), period("2025"), derived=False) + assert sequence_number( + simulation, "independent_input", "2025" + ) >= earliest_source_dependency(simulation) + simulation.set_input(PUBLIC, "2024", [24_000.0]) + np.testing.assert_array_equal(stored(simulation, "independent_input", "2025"), [99]) + assert not holder.is_derived(period("2025")) + assert holder._is_input(period("2025")) is registered + + +def test_reviewer_77_probe_has_no_source_dependency_and_keeps_fast_cache(): + simulation = fresh() + assert scalar(simulation, "independent_formula", "2024") == 77 + value = stored(simulation, "independent_formula", "2024") + fast_cache = dict(simulation._fast_cache) + assert earliest_source_dependency(simulation) is None + epoch = simulation._input_epoch + simulation.set_input(PUBLIC, "2024", [12_000.0]) + # This is the reviewer's 77 probe: no source has been read, so there is + # no since and the adapter must drop nothing, including fast-cache data. + assert stored(simulation, "independent_formula", "2024") is value + assert simulation._input_epoch == epoch + for key, cached in fast_cache.items(): + assert simulation._fast_cache[key] is cached + + +def test_reload_reinstall_write_and_read(): + simulation = fresh({PUBLIC: {"2024": [12_000.0]}}) + importlib.reload(hooks) + hooks.install_social_security_input_compatibility() + simulation.set_input(PUBLIC, "2024", [24_000.0]) + assert scalar(simulation, PUBLIC + "_reported", "2024") == 24_000 + assert scalar(simulation, PUBLIC, "2024") == 24_000 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-q"])) diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index c319a8cc5eb..8cfc229e005 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -280,7 +280,6 @@ def migrate_social_security_inputs( # no successful holder writes to route, including in constructor probes. if getattr(simulation, "persons", None) is None: return moved - purged = False for name in SOCIAL_SECURITY_COMPONENTS: holder = _holder(simulation, name) if holder is None: @@ -305,14 +304,7 @@ def migrate_social_security_inputs( # Routing removes an exact public override and exposes its # formula. A custom input handler may already have calculated # values from that override while the builder was loading. - simulation._drop_values_that_may_depend_on(name, period) - if not purged: - # Core's dependency drops apply only on branches. A - # builder can replace a source on an ordinary simulation - # whose base or downstream values are already cached. - # Keep raw carry seeds and every supplied input intact. - _purge_derived_arrays(simulation) - purged = True + _purge_derived_arrays(simulation, name, period) _store_source(holder, period, value.array, source_branch) if explicit is None: # A local source must also displace inherited public @@ -479,12 +471,24 @@ def _select_subsample_sources(simulation, metadata, person_ids): } -def _purge_derived_arrays(simulation): - """Discard formula results after benefit sources change. +def _purge_derived_arrays(simulation, public, period): + """Drop only derived arrays Core numbers as possibly dependent. - Core's general cache drop also removes nonregistered, nonderived carry - seeds. Those are independent caller values and remain valid here. + Refresh-scope ruling (2026-10-26), rules 1–4: use the earliest dependency + of either source name, retain raw values and supplied inputs, and do + nothing without a dependency record. Core's general _drop_computed also + removes unregistered nonderived seeds, so it cannot serve this adapter. """ + history = simulation._get_store_history() + dependencies = [ + since + for name in (public, f"{public}_reported") + if (since := history.earliest_dependency(name, to_period(period))) is not None + ] + if not dependencies: + return 0 + since = min(dependencies) + dropped = 0 for population in simulation.populations.values(): for holder in population._holders.values(): for tier, storage in ( @@ -494,18 +498,52 @@ def _purge_derived_arrays(simulation): if storage is None: continue for branch, period in tuple(storage.get_known_branch_periods()): - if storage.is_derived(period, branch): + key = ( + f"{branch}:{period}" + if tier == "memory" + else f"{branch}_{period}" + ) + number = storage._sequence_numbers.get(key) + if ( + storage.is_derived(period, branch) + and number is not None + and number >= since + ): _delete_exact(holder, period, branch, tier=tier) + dropped += 1 + if not dropped: + return 0 + # _delete_exact uses Holder._evict_fast_cache per removed key. Also drop + # cached returns whose holders do not keep arrays (e.g. cache blacklist), + # just as Core's _drop_computed does, but only after an actual deletion. simulation._fast_cache = {} + from policyengine_core.simulations.simulation import ( + _calculation_frames, + _taint_waiting_frames, + ) + + simulation._input_epoch += 1 + frames = _calculation_frames.get() + if not frames: + _taint_waiting_frames(simulation) + for frame in frames: + frame.drop_epochs.setdefault(simulation, set()).add(simulation._input_epoch) + # Do not prune history: raw seeds and their readers can remain, unlike + # after Core's broader _drop_computed. Their dependency records still hold. + return dropped def install_social_security_input_compatibility(): from policyengine_core.holders import Holder + from policyengine_core.holders import helpers from policyengine_core.simulations import Simulation as CoreSimulation from policyengine_core.simulations.simulation_builder import SimulationBuilder from policyengine_core.tools import simulation_dumper if getattr(Holder.get_array, "_social_security_readtime", False): + # importlib.reload replaces this module's globals while the installed + # wrappers remain on Core classes. Recover their physical originals. + _PHYSICAL.update(Holder.get_array._social_security_physical) return if getattr(Holder._set, "_social_security_compatibility", False): raise RuntimeError( @@ -520,8 +558,33 @@ def install_social_security_input_compatibility(): recorded_periods=CoreSimulation._get_set_input_periods, recorded_array=CoreSimulation._get_recorded_input_array, calculate=CoreSimulation._calculate, + stored_input=helpers.get_stored_input, ) + @wraps(_PHYSICAL["stored_input"]) + def get_stored_input(holder, period, branch_name): + value = _PHYSICAL["stored_input"](holder, period, branch_name) + if ( + value is None + and holder.variable.name in SOCIAL_SECURITY_COMPONENTS + and _enabled(holder.simulation) + ): + # Core's divide/dispatch helpers must count an initial input even + # after routing hides it from public formula reads. Match Core's + # exact input branch and supplied-input tier checks; raw seeds and + # formula caches do not reserve a helper's subperiod allocation. + source = _physical_value( + holder, + period, + (f"{_SOURCE_BRANCH}{branch_name}",), + registered_only=True, + ) + if source is not None: + return source.array + return value + + helpers.get_stored_input = get_stored_input + @wraps(_PHYSICAL["get_array"]) def get_array(holder, period, branch_name="default"): simulation, name = holder.simulation, holder.variable.name @@ -541,6 +604,7 @@ def get_array(holder, period, branch_name="default"): return _PHYSICAL["get_array"](holder, period, branch_name) get_array._social_security_readtime = True + get_array._social_security_physical = dict(_PHYSICAL) Holder.get_array = get_array @wraps(_PHYSICAL["get_known_periods"]) @@ -647,11 +711,9 @@ def store( ) if tracked: public = name.removesuffix("_reported") - simulation._drop_values_that_may_depend_on(public, period) - # Component bases and downstream results also depend on the - # reported source. Core's dependency drops run only on branches, - # so ordinary simulations need the same derived-only cleanup. - _purge_derived_arrays(simulation) + # Holder.set_input keeps Core's own branch dependency drop. A + # low-level _set has no such Core drop and must retain raw seeds. + _purge_derived_arrays(simulation, public, period) result = _PHYSICAL["set"]( holder, period, From 2e7cccd79930b4b7f36cdd1d62e30b02be2c47db Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 03:39:35 -0700 Subject: [PATCH 52/76] Compare compound benefit helpers with pinned current law --- .../test_social_security_current_law.py | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index c103a29f4e1..138254d09e3 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -28,6 +28,12 @@ "social_security_survivors", "social_security_dependents", ) +COMPOUND_HELPER_ROUTES = tuple( + f"compound_{helper}_{component}_{source}" + for component in COMPONENTS + for helper in ("divide", "dispatch", "ordinary_month") + for source in ("empty", "initial") +) BENEFITS = (*COMPONENTS, "social_security") DOWNSTREAM = ( "social_security_taxable_self_employment_income", @@ -171,6 +177,10 @@ def _probe(): import pandas as pd import yaml from policyengine_core.data import Dataset + from policyengine_core.holders.helpers import ( + set_input_dispatch_by_period, + set_input_divide_by_period, + ) from policyengine_core.periods import period from policyengine_core.reforms import Reform from policyengine_core.simulations import ( @@ -385,6 +395,38 @@ def record_live_value(route, simulation, variable, year): del simulation gc.collect() + # Core's compound helpers retain each existing annual source before + # allocating the remainder. Initial ingress must stay visible to that + # allocation even when the adapter reserves its physical source slot. + # No normalization is used: these are ordinary set-then-calculate routes + # and must equal unmodified main without any intended difference. + for component in COMPONENTS: + for helper_name, helper in ( + ("divide", set_input_divide_by_period), + ("dispatch", set_input_dispatch_by_period), + ("ordinary_month", None), + ): + for source in ("empty", "initial"): + inputs = extraction_situation() + person_inputs = inputs["people"]["retiree"] + del person_inputs["social_security_retirement"] + if source == "initial": + person_inputs[component] = {2024: 12_000} + simulation = SimulationBuilder().build_from_dict(system, inputs) + if helper is None: + simulation.set_input(component, "month:2024-01:24", [36_000]) + else: + helper( + simulation.get_holder(component), + period("year:2024:2"), + np.array([36_000]), + ) + result[f"compound_{helper_name}_{component}_{source}"] = _snapshot( + simulation, (2024, 2025, 2026), (component, "social_security") + ) + del simulation + gc.collect() + simulation = SimulationBuilder().build_from_entities(system, situation()) result["core_entities_builder"] = _snapshot(simulation, (2024, 2026)) del simulation @@ -1022,6 +1064,7 @@ def differential(tmp_path_factory): "multi_year_simulation", "multi_year_microsimulation", "h5_scored_year", + *COMPOUND_HELPER_ROUTES, ), ) def test_legacy_public_ingress_matches_unmodified_main(differential, route): From a439daee9581f834a1adc710b643dadfb72151b4 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 04:12:29 -0700 Subject: [PATCH 53/76] Refresh rebuilt and exact source caches within dependency bounds --- .../core/test_social_security_real_core.py | 112 +++++++++++- .../tools/social_security_inputs.py | 159 +++++++++++++++--- 2 files changed, 244 insertions(+), 27 deletions(-) diff --git a/policyengine_us/tests/core/test_social_security_real_core.py b/policyengine_us/tests/core/test_social_security_real_core.py index b266e7a0a78..9825d821d00 100644 --- a/policyengine_us/tests/core/test_social_security_real_core.py +++ b/policyengine_us/tests/core/test_social_security_real_core.py @@ -345,5 +345,115 @@ def test_reload_reinstall_write_and_read(): assert scalar(simulation, PUBLIC, "2024") == 24_000 +@pytest.mark.parametrize("warm_source", (False, True)) +def test_dataset_rebuild_refreshes_discarded_holder_caches(warm_source): + from policyengine_core.data import Dataset + + simulation = fresh() + assert scalar(simulation, "independent_formula", "2023") == 77 + independent_holder = simulation.get_holder("independent_formula") + independent_key = ("independent_formula", period("2023")) + independent_fast = simulation._fast_cache[independent_key] + if warm_source: + assert scalar(simulation, PUBLIC, "2023") == 0 + + class TinyDataset: + data_format = Dataset.TIME_PERIOD_ARRAYS + time_period = "2023" + + def load(self): + return { + "person_id": {"2023": np.array([1])}, + "household_id": {"2023": np.array([1])}, + "person_household_id": {"2023": np.array([1])}, + PUBLIC: {"2023": np.array([1.0])}, + } + + simulation.dataset = TinyDataset() + simulation.build_from_dataset() + assert scalar(simulation, PUBLIC + "_reported", "2023") == 1 + assert scalar(simulation, PUBLIC, "2023") == 1 + # Core replaces populations during dataset loading. Retain its former + # holder here to audit the adapter's treatment of the detached cache. + assert float(independent_holder.get_array(period("2023"))[0]) == 77 + if not warm_source: + # No source read means no cutoff and no adapter eviction, including + # a cached reader whose former holder Core itself has detached. + assert simulation._fast_cache[independent_key] is independent_fast + + +def test_dataset_rebuild_keeps_new_input_provenance_when_old_cache_shares_key(): + from policyengine_core.data import Dataset + + simulation = SimulationBuilder().build_from_variables(make_system(), {}) + assert scalar(simulation, PUBLIC, "2023") == 0 + since = earliest_source_dependency(simulation, requested="2023") + assert since is not None + assert scalar(simulation, "independent_input", "2023") == 0 + old_holder = simulation.get_holder("independent_input") + assert old_holder.is_derived(period("2023")) + assert sequence_number(simulation, "independent_input", "2023") >= since + + class TinyDataset: + data_format = Dataset.TIME_PERIOD_ARRAYS + time_period = "2023" + + def load(self): + return { + "person_id": {"2023": np.array([1])}, + "household_id": {"2023": np.array([1])}, + "person_household_id": {"2023": np.array([1])}, + # Store this before the benefit write so an old derived + # cache with the same key cannot unregister the new input. + "independent_input": {"2023": np.array([99.0])}, + PUBLIC: {"2023": np.array([1.0])}, + } + + simulation.dataset = TinyDataset() + simulation.build_from_dataset() + current_holder = simulation.get_holder("independent_input") + assert current_holder is not old_holder + np.testing.assert_array_equal(stored(simulation, "independent_input", "2023"), [99]) + assert current_holder._is_input(period("2023")) + assert period("2023") in simulation._get_set_input_periods("independent_input") + snapshot = hooks.export_source_inputs(simulation) + person = next(iter(snapshot["sources"]["people"].values())) + assert person["independent_input"]["2023"] == 99 + assert scalar(simulation, PUBLIC, "2023") == 1 + + +@pytest.mark.parametrize("initial", (False, True), ids=("raw_source", "initial_source")) +def test_public_write_evicts_exact_reported_memo_without_dropping_other_caches(initial): + simulation = fresh() + assert scalar(simulation, "independent_formula", "2024") == 77 + independent_key = ("independent_formula", period("2024")) + independent_fast = simulation._fast_cache[independent_key] + if initial: + SimulationBuilder().build_from_variables( + simulation.tax_benefit_system, + {PUBLIC: {"2024": [12_000.0]}}, + simulation=simulation, + ) + else: + simulation.get_holder(PUBLIC).put_in_cache( + np.array([12_000.0]), period("2024"), derived=False + ) + assert scalar(simulation, PUBLIC + "_reported", "2024") == 12_000 + # This exact read returns the source view without storing a derived + # reported array. There is no post-cutoff formula array to purge. + assert ( + simulation.get_holder(PUBLIC + "_reported")._memory_storage.get(period("2024")) + is None + ) + epoch = simulation._input_epoch + simulation.set_input(PUBLIC, "2024", [24_000.0]) + assert scalar(simulation, PUBLIC + "_reported", "2024") == 24_000 + assert simulation._fast_cache[independent_key] is independent_fast + assert stored(simulation, "independent_formula", "2024") is independent_fast + assert simulation._input_epoch == epoch + + if __name__ == "__main__": - raise SystemExit(pytest.main([__file__, "-q"])) + raise SystemExit( + pytest.main([__file__, "-q", "--confcutdir", str(Path(__file__).parent)]) + ) diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 8cfc229e005..e2255d01e32 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -18,6 +18,7 @@ from dataclasses import dataclass from functools import wraps from pathlib import Path +from weakref import ref from policyengine_core.data_storage.store_history import next_sequence_number from policyengine_core.periods import period as to_period @@ -40,6 +41,9 @@ _initial_social_security_ingress = ContextVar("initial_ss_ingress", default=False) _initial_social_security_writes = ContextVar("initial_ss_writes", default=None) +_loading_social_security_cache_entries = ContextVar( + "loading_ss_cache_entries", default=() +) _reported_request = ContextVar("reported_ss_request", default=None) _subsampling_social_security_sources = ContextVar( "subsampling_ss_sources", default=None @@ -471,6 +475,76 @@ def _select_subsample_sources(simulation, metadata, person_ids): } +def _derived_cached_returns(simulation): + """Remember numbered formula returns without retaining an old dataset. + + Core's dataset rebuild detaches holders but leaves their fast-cache + returns. Weak references and sequence numbers let the bounded refresh + remove those returns after a source write, without retaining populations + or using detached holders to modify the new input registry. + """ + fast_cache = getattr(simulation, "_fast_cache", {}) + if not fast_cache: + return {} + holders = { + holder.variable.name: holder + for population in simulation.populations.values() + for holder in population._holders.values() + } + + def has_source(holder, period, branches): + return holder is not None and any( + storage is not None + and storage.has(period, branch) + and not storage.is_derived(period, branch) + for branch in branches + for storage in (holder._memory_storage, holder._disk_storage) + ) + + entries = {} + for key, array in fast_cache.items(): + name, requested = key + holder = holders.get(name) + if holder is None: + continue + period = holder._storage_period(requested) + if name in _REPORTED_COMPONENTS and not holder.variable.is_neutralized: + branches = _branches(simulation) + # A read-time explicit/public source can shadow a physical + # reported cache. Such a returned input is not a derived array. + if has_source(holder, period, branches): + continue + exact_only = period.unit == "year" and period.size > 1 + explicit_before = any( + (known == period if exact_only else known.start <= period.start) + and has_source(holder, known, branches) + for known in _physical_periods(holder) + ) + if not explicit_before and has_source( + holders.get(name.removesuffix("_reported")), + period, + _source_branches(simulation), + ): + continue + branch = holder._branch_storing(period, simulation.branch_name) + if branch is None: + continue + if holder._memory_storage.has(period, branch): + storage = holder._memory_storage + storage_key = f"{branch}:{period}" + else: + storage = holder._disk_storage + storage_key = f"{branch}_{period}" + if storage.is_derived(period, branch): + number = storage._sequence_numbers.get(storage_key) + if number is not None: + # Do not read disk arrays: decoding produces a new object. + # The fast-cache return already has the selected slot's + # derivation/sequence metadata; retain only a weak reference. + entries[key] = (number, ref(array)) + return entries + + def _purge_derived_arrays(simulation, public, period): """Drop only derived arrays Core numbers as possibly dependent. @@ -511,6 +585,18 @@ def _purge_derived_arrays(simulation, public, period): ): _delete_exact(holder, period, branch, tier=tier) dropped += 1 + for owner, entries in _loading_social_security_cache_entries.get(): + if owner is not simulation: + continue + for key, (number, reference) in tuple(entries.items()): + if number >= since: + array = reference() + if array is not None and simulation._fast_cache.get(key) is array: + # This numbered formula array remains in the fast cache + # after Core detached its original holder during rebuild. + simulation._fast_cache.pop(key) + dropped += 1 + entries.pop(key) if not dropped: return 0 # _delete_exact uses Holder._evict_fast_cache per removed key. Also drop @@ -732,6 +818,16 @@ def store( if registered and contexts and branch_name == "default" else branch_name ) + if name in SOCIAL_SECURITY_COMPONENTS: + reported_holder = _holder(simulation, f"{name}_reported") + if reported_holder is not None: + # An exact reported read can memoize the public input + # itself, without storing any derived array. Evict that + # changed source view exactly as Core evicts the written + # holder's own memo; unrelated cached returns stay intact. + reported_holder._evict_fast_cache( + holder._storage_period(period), effective_branch + ) writes = _initial_social_security_writes.get() if name in SOCIAL_SECURITY_COMPONENTS and registered and writes is not None: writes.setdefault(simulation, set()).add( @@ -782,6 +878,14 @@ def load(first, *args, **kwargs): if writes is None: writes = {} writes_token = _initial_social_security_writes.set(writes) + cache_entries_token = _loading_social_security_cache_entries.set( + ( + *_loading_social_security_cache_entries.get(), + (first, _derived_cached_returns(first)), + ) + if dataset + else _loading_social_security_cache_entries.get() + ) result = None try: result = original(first, *args, **kwargs) @@ -789,33 +893,36 @@ def load(first, *args, **kwargs): finally: _initial_social_security_writes.reset(writes_token) _initial_social_security_ingress.reset(token) - if outermost: - targets = dict(writes) - target = first if dataset else result - if target is None and not dataset: - target = kwargs.get("simulation") - if target is None and len(args) > 2: - target = args[2] - if target is not None: - targets.setdefault(target, set()) - for target, loaded in targets.items(): - # During routing, projected twins temporarily retain - # any both-name public final overrides. Conversion - # then replaces those projection arrays with the - # independently serialized public source data. - migrate_social_security_inputs( - target, loaded_public_inputs=loaded - ) - rebuilding = _subsampling_social_security_sources.get() - if ( - dataset - and rebuilding is not None - and rebuilding[0] is target - ): - metadata = _select_subsample_sources( - target, rebuilding[1], rebuilding[2] + try: + if outermost: + targets = dict(writes) + target = first if dataset else result + if target is None and not dataset: + target = kwargs.get("simulation") + if target is None and len(args) > 2: + target = args[2] + if target is not None: + targets.setdefault(target, set()) + for target, loaded in targets.items(): + # During routing, projected twins temporarily retain + # any both-name public final overrides. Conversion + # then replaces those projection arrays with the + # independently serialized public source data. + migrate_social_security_inputs( + target, loaded_public_inputs=loaded ) - _restore_automatic_reported_metadata(target, metadata) + rebuilding = _subsampling_social_security_sources.get() + if ( + dataset + and rebuilding is not None + and rebuilding[0] is target + ): + metadata = _select_subsample_sources( + target, rebuilding[1], rebuilding[2] + ) + _restore_automatic_reported_metadata(target, metadata) + finally: + _loading_social_security_cache_entries.reset(cache_entries_token) return load From 82d52dbfeaad5d924b76a180c5d2966d80d12500 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 04:17:33 -0700 Subject: [PATCH 54/76] Exercise real Core bookkeeping in benefit input state machine --- .../core/_social_security_state_backend.py | 355 ++++++------------ ...est_social_security_input_state_machine.py | 235 ++++++++++++ 2 files changed, 344 insertions(+), 246 deletions(-) diff --git a/policyengine_us/tests/core/_social_security_state_backend.py b/policyengine_us/tests/core/_social_security_state_backend.py index 4ff5caed4c5..706f469db6a 100644 --- a/policyengine_us/tests/core/_social_security_state_backend.py +++ b/policyengine_us/tests/core/_social_security_state_backend.py @@ -2,26 +2,47 @@ The country is intentionally never imported. Core's installed source bodies implement every write, builder, copy, replay, extraction and dump route below; -the US module's complete function bodies implement the hooks. Only entities, -arrays, periods, storage and dependency bookkeeping are synthetic. This tests -input semantics and Core's source selection, not country benefit formulas or -full-model performance. A reported -> base -> public formula chain exercises +the US module's complete function bodies implement the hooks. Entity and +variable descriptors, dataset construction, and disk dump bytes are synthetic. +Periods, memory storage, store history, sequence numbers, dependency drops, +eviction and period helpers use Core itself. This tests input semantics and +Core's source selection, not country benefit formulas or full-model performance. A reported -> base -> public formula chain exercises cache invalidation and scored-year factors. The synthetic uprating index doubles each year so source changes have exact reader values. + +Remaining Core substitutions are limited to fixture plumbing: +- Holder.__init__ and default_array construct one-person float arrays in + memory; create_disk_storage supplies the JSON-backed Disk used only by + dump/restore. +- Population construction, get_holder and clone provide one person and one + household. System descriptors and parameters replace country legislation. +- Simulation.__init__, build_from_populations, get_holder and + get_variable_population attach those populations. calculate provides its + single-thread fast-cache entry point, _run_formula defines the tiny + reported/base/public chain, _cast_formula_result converts float arrays, + _check_period_consistency accepts exact annual reads, and _check_for_cycle + skips tracing for this acyclic chain. check_macro_cache returns False; + subsample is unsupported because the fixture contains one person. +- Builder entity registration/declaration and membership methods only create + that one-person situation; _get_person_count and check_type reflect its + fixed shape. Dataset.load supplies the tiny dictionary input payload. + get_parameter returns the fixture's doubling uprating index. +- Dumper _dump_entity/_restore_entity replace entity serialization only; + Disk replaces OnDiskStorage's physical bytes. Public dump/restore bodies, + holder provenance, memory tiers and source restoration remain Core/US code. """ from __future__ import annotations import ast import copy -import itertools import json import logging import os import sys import types import warnings -from contextlib import ExitStack +from contextlib import ExitStack, nullcontext from contextvars import ContextVar from dataclasses import dataclass from enum import Enum @@ -29,8 +50,21 @@ from importlib.metadata import distribution from pathlib import Path from unittest.mock import patch +from weakref import ref import numpy as np +from policyengine_core import commons, periods +from policyengine_core.data_storage.in_memory_storage import InMemoryStorage +from policyengine_core.data_storage.store_history import ( + StoreHistory, + _period_from_storage_key, + next_sequence_number, +) +from policyengine_core.periods import Period +from policyengine_core.simulations import simulation as core_simulation + + +as_period = periods.period ROOT = Path(__file__).resolve().parents[2] @@ -42,55 +76,6 @@ ) -@dataclass(frozen=True, order=True) -class Instant: - year: int - month: int = 1 - - @property - def date(self): - return (self.year, self.month) - - -@dataclass(frozen=True) -class Period: - start: Instant - size: int = 1 - unit: str = "year" - - def __iter__(self): - return iter((self.unit, (self.start.year, self.start.month, 1), self.size)) - - @property - def stop(self): - return Instant(self.start.year + self.size - 1, 12) - - def __str__(self): - if self.unit == "eternity": - return "eternity" - return ( - str(self.start.year) - if self.size == 1 - else f"year:{self.start.year}:{self.size}" - ) - - def contains(self, other): - return self.start <= other.start and self.stop >= other.stop - - -def as_period(value): - if isinstance(value, Period): - return value - if str(value) == "eternity": - return Period(Instant(0), unit="eternity") - parts = str(value).split(":") - return ( - Period(Instant(int(parts[-2])), int(parts[-1])) - if len(parts) == 3 - else Period(Instant(int(value))) - ) - - @lru_cache(maxsize=None) def source_tree(source): return ast.parse(source.read_text()) @@ -147,91 +132,6 @@ def compiled_hooks(source): return compile(code, str(source), "exec") -class Storage: - """Memory arrays plus Core's separate carry-over and input-tier marks.""" - - is_eternal = False - - def __init__(self): - self._arrays = {} - self._inputs = set() - self._sequence_numbers = {} - - def get(self, period, branch_name="default"): - return self._arrays.get(f"{branch_name}:{as_period(period)}") - - def has(self, period, branch_name="default"): - return f"{branch_name}:{as_period(period)}" in self._arrays - - def is_derived(self, period, branch_name="default"): - return ( - self.has(period, branch_name) - and f"{branch_name}:{as_period(period)}" not in self._inputs - ) - - def put( - self, value, period, branch_name="default", derived=False, sequence_number=None - ): - key = f"{branch_name}:{as_period(period)}" - self._arrays[key] = np.asarray(value).copy() - self._sequence_numbers[key] = sequence_number - if derived: - self._inputs.discard(key) - else: - self._inputs.add(key) - - def get_known_branch_periods(self): - return [ - (key.split(":", 1)[0], as_period(key.split(":", 1)[1])) - for key in self._arrays - ] - - def get_known_periods(self): - return list( - dict.fromkeys(period for _, period in self.get_known_branch_periods()) - ) - - def delete(self, period=None, branch_name="default"): - for key in list(self._arrays): - branch, when = key.split(":", 1) - if branch == branch_name and ( - period is None or as_period(period).contains(as_period(when)) - ): - del self._arrays[key] - self._unmark_dropped_keys() - - def mark_derived_except(self, kept): - self._inputs.intersection_update(kept) - - def drop_computed(self, since=None): - dropped = [ - key - for key in self._arrays - if key not in self._inputs - and (since is None or (self._sequence_numbers.get(key) or 0) >= since) - ] - for key in dropped: - del self._arrays[key] - self._unmark_dropped_keys() - return len(dropped) - - def _unmark_dropped_keys(self): - self._inputs.intersection_update(self._arrays) - - def _stop_sharing_dropped_keys(self): - pass - - def _forget_dropped_numbers(self): - self._sequence_numbers = { - key: number - for key, number in self._sequence_numbers.items() - if key in self._arrays - } - - def clone(self, **kwargs): - return copy.deepcopy(self) - - class Disk: """Tiny on-disk adapter used by unchanged public Core dump functions.""" @@ -260,20 +160,6 @@ def get(self, period): return self.values[as_period(period)] -class History: - def copy(self): - return History() - - def record_unknown_sources(self, *args): - pass - - def record_derived(self, *args): - pass - - def prune(self, *args): - pass - - class Backend: """Install the unchanged hooks around unchanged Core routes in isolation.""" @@ -287,47 +173,10 @@ def __init__(self): def _setup(self): core = Path(distribution("policyengine-core").locate_file("policyengine_core")) - counter = itertools.count(1) - - def next_sequence(): - return next(counter) - - periods = types.SimpleNamespace( - period=as_period, - YEAR="year", - MONTH="month", - ETERNITY="eternity", - DAY="day", - key_period_size=lambda p: p.size, - unit_weight=lambda unit: {"day": 1, "month": 2, "year": 3}[unit], - ) - commons = types.SimpleNamespace( - empty_clone=lambda value: object.__new__(type(value)) - ) - calculation_frames = ContextVar("state_machine_frames", default=()) - formula_simulation = ContextVar("state_machine_formula", default=None) - branch_clone = ContextVar("state_machine_branch_clone", default=None) - - class BranchClone: - def __init__(self, simulation): - self.simulation = simulation - self.pending = True - - class Encoder(json.JSONEncoder): - def default(self, value): - if isinstance(value, np.ndarray): - return value.tolist() - if isinstance(value, np.generic): - return value.item() - return super().default(value) - - @dataclass - class PreservedInput: - variable_name: str - branch_name: str - period: Period - value: object - storage: str + calculation_frames = core_simulation._calculation_frames + formula_simulation = core_simulation._formula_simulation + branch_clone = core_simulation._branch_clone + BranchClone = core_simulation._BranchClone person = types.SimpleNamespace(key="person", plural="people", is_person=True) member = types.SimpleNamespace( @@ -346,7 +195,7 @@ def __init__(self, population, name): self.population = population self.simulation = population.simulation self.variable = self.simulation.tax_benefit_system.get_variable(name) - self._memory_storage = Storage() + self._memory_storage = InMemoryStorage(is_eternal=False) self._disk_storage = None self._on_disk_storable = False self._do_not_store = False @@ -357,27 +206,6 @@ def _storage_period(self, period): def default_array(self): return np.zeros(1) - def _raise_if_input_contains_nan(self, array): - assert np.isfinite(array).all() - - def _to_array(self, value, validate_nan=False): - return np.asarray(value, dtype=float).reshape(1) - - def _record_store(self, *args): - pass - - def _record_inputs(self, *args): - pass - - def _evict_fast_cache(self, *args, **kwargs): - pass - - def _input_unchanged(self, *args): - return False - - def _stores(self, period, branch_name): - return self._memory_storage.has(period, branch_name) - def _get_array_from_storage(self, period, branch_name="default"): return self._memory_storage.get(period, branch_name) @@ -439,6 +267,7 @@ def __init__(self): definition_period="year", is_neutralized=False, value_type=float, + dtype=np.dtype("float64"), set_input=None, end=None, requires_computation_after=None, @@ -447,9 +276,22 @@ def __init__(self): calculate_output=None, quantity_type="flow", check_uprating_value_type=lambda: None, - is_input_variable=lambda name=name: name.endswith("_reported"), + is_input_variable=lambda name=name: ( + name.endswith("_reported") + or name + in ( + "independent_input", + "raw_seed", + "early_derived", + "late_derived", + ) + ), ) for name in ( + "independent_input", + "raw_seed", + "early_derived", + "late_derived", *COMPONENTS, *(f"{name}_reported" for name in COMPONENTS), *(f"{name}_base" for name in COMPONENTS), @@ -476,6 +318,7 @@ def parameters(self, *args): system = System() class Simulation: + _open_frames = None default_role = "member" default_input_period = "2024" default_calculation_period = "2024" @@ -486,6 +329,8 @@ def __init__(self, tax_benefit_system=None, populations=None, **kwargs): self._user_input_contexts = [] self._inputs_set = 0 self._input_epoch = 0 + self._store_history = None + self._requested_variables = None self._calculations_started = 0 self._calculations_in_flight = 0 self.input_variables = [] @@ -522,18 +367,6 @@ def get_holder(self, name): def get_variable_population(self, name): return self.persons - def _drop_values_that_may_depend_on(self, *args): - pass - - def _get_store_history(self): - return History() - - def _get_requested_variables(self): - return set() - - def _drop_computed(self, since=None): - raise AssertionError("Installed Core cache cleanup was not loaded") - def subsample(self, *args): raise NotImplementedError( "Dataset subsampling is outside this storage backend" @@ -651,20 +484,25 @@ def load(self): logging=logging, log=logging.getLogger(__name__), commons=commons, - next_sequence_number=next_sequence, + next_sequence_number=next_sequence_number, _calculation_frames=calculation_frames, _formula_simulation=formula_simulation, _branch_clone=branch_clone, _BranchClone=BranchClone, - _note_unobserved_activity=lambda: None, - _taint_waiting_frames=lambda *args: None, + _note_unobserved_activity=core_simulation._note_unobserved_activity, + _taint_waiting_frames=core_simulation._taint_waiting_frames, + _frame_for=core_simulation._frame_for, + _RERUNS_AFTER_INPUT_CHANGE=core_simulation._RERUNS_AFTER_INPUT_CHANGE, + nullcontext=nullcontext, directory_containing=lambda path: None, - PreservedUserInput=PreservedInput, - NpEncoder=Encoder, + PreservedUserInput=core_simulation.PreservedUserInput, + StoreHistory=StoreHistory, + NpEncoder=core_simulation.NpEncoder, ETERNITY="eternity", YEAR="year", MONTH="month", QuantityType=types.SimpleNamespace(STOCK="stock"), + Enum=Enum, SpiralError=type("SpiralError", (Exception,), {}), FullTracer=type("FullTracer", (), {}), get_parameter=lambda *args: lambda instant: 2 ** (instant.year - 2023), @@ -674,7 +512,7 @@ def load(self): _get_person_count=lambda values: 1, check_type=lambda *args: None, PeriodMismatchError=ValueError, - _period_from_storage_key=as_period, + _period_from_storage_key=_period_from_storage_key, copy=copy, ) holder_methods = ( @@ -685,6 +523,17 @@ def load(self): "get_known_branch_periods", "set_input", "_set", + "_to_array", + "_raise_if_input_contains_nan", + "_record_store", + "_record_inputs", + "_evict_fast_cache", + "_input_unchanged", + "_is_input", + "_has_unnumbered_values", + "_stored_sequence_number", + "_storage_period", + "_stores", "put_in_cache", "_seed_input_storage", "_record_input_storage", @@ -711,7 +560,15 @@ def load(self): for node in source_tree(helper_source).body if isinstance(node, ast.FunctionDef) ) - core_helpers = execute(helper_source, helper_names, namespace) + helpers = types.ModuleType("policyengine_core.holders.helpers") + helpers.__dict__.update(namespace) + helpers.__dict__["__name__"] = helpers.__name__ + helpers.__dict__["__package__"] = "policyengine_core.holders" + core_helpers = execute(helper_source, helper_names, helpers.__dict__) + namespace.update(core_helpers) + # Variable supplies this default handler for yearly FLOW variables. + for variable in system.variables.values(): + variable.set_input = core_helpers["set_input_divide_by_period"] # Core 3.33.12 cleans overlapping raw seeds after a direct set_input; # 3.33.2 does this only in dispatch/divide helpers. Keep the reference # model tied to the installed behavior rather than a version string. @@ -737,6 +594,9 @@ def load(self): ) sim_methods = ( "set_input", + "_get_store_history", + "_get_requested_variables", + "_drop_values_that_may_depend_on", "build_from_dataset", "clone", "get_branch", @@ -748,6 +608,8 @@ def load(self): "_get_set_input_periods", "_get_recorded_input_array", "_calculate", + "calculate_add", + "_cache_option_result", "_calculation_start", "_cache_result", "_may_keep", @@ -785,7 +647,7 @@ def module(name, **values): Simulation=Simulation, OnDiskStorage=Disk, ETERNITY="eternity", - next_sequence_number=next_sequence, + next_sequence_number=next_sequence_number, DERIVED_PERIODS_FILE="derived_periods.txt", DUMP_FORMAT_FILE="dump_format.txt", DUMP_FORMAT_VERSION=1, @@ -803,11 +665,6 @@ def module(name, **values): ), dumper.__dict__, ) - helpers = module( - "policyengine_core.holders.helpers", - _period_from_storage_key=as_period, - **core_helpers, - ) period_hooks = module( "policyengine_us.tools.social_security_periods", np=np, @@ -841,7 +698,10 @@ def module(name, **values): ), "policyengine_core.simulations.simulation": module( "policyengine_core.simulations.simulation", - _note_input_change=lambda *args: None, + _note_input_change=core_simulation._note_input_change, + _calculation_frames=calculation_frames, + _taint_waiting_frames=namespace["_taint_waiting_frames"], + NpEncoder=core_simulation.NpEncoder, ), "policyengine_core.simulations.simulation_builder": module( "policyengine_core.simulations.simulation_builder", @@ -867,13 +727,14 @@ def module(name, **values): "_social_security_state_hooks", json=json, ContextVar=ContextVar, + ref=ref, dataclass=dataclass, wraps=wraps, Path=Path, - next_sequence_number=next_sequence, + next_sequence_number=next_sequence_number, get_stored_input=namespace["get_stored_input"], to_period=as_period, - NpEncoder=Encoder, + NpEncoder=core_simulation.NpEncoder, ETERNITY="eternity", ) self.stack.enter_context( @@ -894,6 +755,8 @@ def module(name, **values): self.system = system self.hooks = hook_namespace self.dumper = dumper + self.helpers = helpers + self.next_sequence_number = next_sequence_number def close(self): self.stack.close() diff --git a/policyengine_us/tests/core/test_social_security_input_state_machine.py b/policyengine_us/tests/core/test_social_security_input_state_machine.py index 768432ee89a..67f3e5dc56f 100644 --- a/policyengine_us/tests/core/test_social_security_input_state_machine.py +++ b/policyengine_us/tests/core/test_social_security_input_state_machine.py @@ -238,6 +238,8 @@ def __init__(self): self.temp = tempfile.TemporaryDirectory(prefix="ss-input-machine-") self.dump_count = 0 self.branch_initial_rewrite_done = False + self.preservation_checks = [] + self.helper_checks = [] def teardown(self): self.temp.cleanup() @@ -268,6 +270,233 @@ def write_key(model, component, period, data): ) return component, period + def prepare_core_input_drop(self, simulation, model, name, period, amount): + """Model Core's own branch drop, including its raw-seed behavior. + + Core first reclassifies all nonregistered values as derived; only + registered inputs remain eligible sources after such a drop. + Low-level adapter writes deliberately never call this method. + """ + if simulation.parent_branch is None: + return + holder = simulation.get_holder(name) + if holder._input_unchanged(as_period(period), [amount], simulation.branch_name): + return + since = simulation._get_store_history().earliest_dependency( + name, as_period(period) + ) + if holder._has_unnumbered_values(): + since = 0 + if since is None: + return + for mapping in (model.public, model.explicit, model.overrides): + for key, value in tuple(mapping.items()): + if not value.registered: + # Core first reclassifies every nonregistered seed as + # derived, even when its sequence is before the cutoff. + # It no longer supplies a source to uprating/carry-over. + del mapping[key] + + @rule( + component=COMPONENT, + helper=st.sampled_from(("divide", "dispatch", "ordinary")), + compound=st.sampled_from(("year:2024:2", "month:2024-01:24", "year:2024-07:2")), + initial=st.booleans(), + amount=st.integers(min_value=24_000, max_value=48_000), + ) + def initial_source_then_period_helper( + self, component, helper, compound, initial, amount + ): + """Compare each allocation with Core running without the adapter.""" + if helper == "ordinary": + compound = "month:2024-01:24" + inputs = {component: {"2024": [12_000]}} if initial else {} + adapted = self.backend.Builder().build_from_variables( + self.backend.system, inputs + ) + self.backend.system._social_security_input_compatibility = False + try: + plain = self.backend.Builder().build_from_variables( + self.backend.system, inputs + ) + self.apply_helper(plain, component, helper, compound, amount) + finally: + self.backend.system._social_security_input_compatibility = True + self.apply_helper(adapted, component, helper, compound, amount) + period = as_period(compound) + branch = adapted.branch_name + # The fixture's public formula multiplies its initial source by two. + # Compare allocations through Core's input lookup, which is the + # source the helper must see, rather than that reform factor. + years = [] + allocation_model = Model() + if initial: + allocation_model.write( + component, + False, + "2024", + 12_000, + True, + "initial helper source", + initial=True, + ) + subperiod = period.start.period("year") + stop = period.start.offset(period.size, period.unit) + while subperiod.start < stop: + expected = self.backend.helpers.get_stored_input( + plain.get_holder(component), subperiod, plain.branch_name + ) + actual = self.backend.helpers.get_stored_input( + adapted.get_holder(component), subperiod, branch + ) + assert expected is not None and actual is not None + self.helper_checks.append( + (adapted, component, subperiod, expected.copy(), compound) + ) + years.append(subperiod) + if not (initial and str(subperiod) == "2024"): + allocation_model.write( + component, + False, + str(subperiod), + float(expected[0]), + True, + "Core helper", + ) + subperiod = subperiod.offset(1) + # Exercise a real compound calculation after allocations. Core's + # calculate_add sums annual subperiod readers instead of fabricating + # an exact total source for the compound. + actual_total = ( + float(adapted.calculate_add(component, period)[0]) + if period.unit == "year" + else sum(float(adapted.calculate(component, year)[0]) for year in years) + ) + calculation_years = ( + period.get_subperiods("year") if period.unit == "year" else years + ) + expected_total = 0 + for year in calculation_years: + override = allocation_model.visible( + allocation_model.overrides, component, str(year) + ) + expected_total += ( + override.amount + if override is not None + else allocation_model.reader_amount(component, str(year)) * 2 + ) + assert actual_total == expected_total + + def apply_helper(self, simulation, component, helper, compound, amount): + if helper == "ordinary": + simulation.set_input(component, compound, [amount]) + else: + function = getattr(self.backend.helpers, f"set_input_{helper}_by_period") + function(simulation.get_holder(component), as_period(compound), [amount]) + + @rule(component=COMPONENT, partly_failing=st.booleans()) + def custom_input_handler_after_initial_source(self, component, partly_failing): + simulation = self.backend.Builder().build_from_variables( + self.backend.system, {component: {"2024": [12_000]}} + ) + variable = self.backend.system.get_variable(component) + original = variable.set_input + failure = RuntimeError("intentional partial input handler failure") + + def handler(holder, period, amount): + if partly_failing: + holder._set(as_period("2025"), [24_000], is_input=True) + # Calculations inside a handler must not preserve old readers + # when the handler subsequently raises after a successful store. + holder.simulation.calculate(component, as_period("2026")) + raise failure + self.backend.helpers.set_input_divide_by_period(holder, period, amount) + + variable.set_input = handler + try: + try: + simulation.set_input(component, "month:2024-01:24", [36_000]) + except RuntimeError as caught: + assert partly_failing and caught is failure + else: + assert not partly_failing + finally: + variable.set_input = original + self.preservation_checks.append( + (simulation, component + "_reported", "2024", 12_000) + ) + self.preservation_checks.append((simulation, component, "2025", 24_000)) + + @rule(component=COMPONENT) + def deleting_explicit_source_reveals_initial_source(self, component): + simulation = self.backend.Builder().build_from_variables( + self.backend.system, {component: {"2024": [12_000]}} + ) + reported = component + "_reported" + simulation.set_input(reported, "2024", [6_000]) + assert float(simulation.calculate(reported, as_period("2024"))[0]) == 6_000 + simulation.get_holder(reported).delete_arrays(as_period("2024")) + # A source read resolves current physical tiers after the deletion. + value = simulation.get_holder(reported).get_array(as_period("2024")) + assert value is not None and float(value[0]) == 12_000 + + @rule(component=COMPONENT, branch=st.booleans(), read_source=st.booleans()) + def bounded_refresh_preserves_caller_values(self, component, branch, read_source): + """Keep pre-cutoff formulas, raw seeds and supplied inputs on every write.""" + simulation = self.backend.DelegatingSimulation() + early = simulation.get_holder("early_derived") + early.put_in_cache([77], as_period("2025"), derived=True) + if read_source: + simulation.set_input(component, "2024", [12_000]) + if branch: + simulation = simulation.get_branch("preservation") + if read_source: + simulation.calculate(component, as_period("2026")) + raw = simulation.get_holder("raw_seed") + raw.put_in_cache([99], as_period("2025"), simulation.branch_name, derived=False) + simulation.set_input("independent_input", "2025", [88]) + late = simulation.get_holder("late_derived") + late.put_in_cache([55], as_period("2025"), simulation.branch_name, derived=True) + since = simulation._get_store_history().earliest_dependency( + component, as_period("2024") + ) + simulation.get_holder(component)._set( + as_period("2024"), [24_000], simulation.branch_name, is_input=True + ) + for name, expected in ( + ("early_derived", 77), + ("raw_seed", 99), + ("independent_input", 88), + ): + self.preservation_checks.append((simulation, name, "2025", expected)) + if since is not None: + # Intended over-approximation: the refresh-scope ruling allows + # an unrelated formula numbered at or after since to be dropped. + assert late.get_array(as_period("2025"), simulation.branch_name) is None + else: + self.preservation_checks.append((simulation, "late_derived", "2025", 55)) + + @invariant() + def caller_values_and_earlier_derived_survive(self): + for simulation, name, period, expected in self.preservation_checks: + value = simulation.get_holder(name).get_array( + as_period(period), simulation.branch_name + ) + assert value is not None and float(value[0]) == expected, ( + name, + period, + "adapter removed a caller value or pre-cutoff formula", + ) + for simulation, component, period, expected, compound in self.helper_checks: + actual = self.backend.helpers.get_stored_input( + simulation.get_holder(component), period, simulation.branch_name + ) + assert actual is not None and (actual == expected).all(), ( + component, + compound, + "Core helper allocation", + ) + @rule( slot=TARGET, component=COMPONENT, @@ -281,6 +510,7 @@ def set_input(self, slot, component, reported, period, amount, route, data=None) simulation, model = self.state(slot) component, period = self.write_key(model, component, period, data) name = self.name(component, reported) + self.prepare_core_input_drop(simulation, model, name, period, amount) if route == "holder": simulation.get_holder(name).set_input( as_period(period), [amount], simulation.branch_name @@ -439,6 +669,8 @@ def builder_load( simulation, model = None, Model() name = self.name(component, reported) builder = self.backend.Builder() + if simulation is not None: + self.prepare_core_input_drop(simulation, model, name, period, amount) if yaml: # This executes Core build_from_dict -> build_from_entities -> # finalize_variables_init -> Holder.set_input, including nested @@ -490,6 +722,9 @@ def setter_stores_then_raises( simulation.__class__ = self.backend.DelegatingSimulation simulation.fail_after_store = True simulation.fail_name = self.name(component, reported) + self.prepare_core_input_drop( + simulation, model, simulation.fail_name, period, amount + ) error = None try: name = self.name(component, reported) From 57dc1952aeb1b92fd2db8db831cbc5a19f7914d1 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 08:54:55 -0700 Subject: [PATCH 55/76] Isolate synthetic Social Security checks from microsimulation CI Run the unchanged input checks in a dedicated Rest group and preserve three isolated processes, distinct reports, and failure propagation. Avoid the differential parent's redundant country import and allow slower probes within a bounded hang guard. --- Makefile | 21 +++++--- .../code_health/test_rest_python_groups.py | 48 ++++++++++++++----- .../test_social_security_inputs.py | 0 .../test_social_security_current_law.py | 4 +- 4 files changed, 54 insertions(+), 19 deletions(-) rename policyengine_us/tests/{microsimulation => core}/test_social_security_inputs.py (100%) diff --git a/Makefile b/Makefile index 12dc55c4b12..2efebc7c580 100644 --- a/Makefile +++ b/Makefile @@ -52,7 +52,7 @@ REST_SPM_TESTS := $(TESTS)/core/test_spm_policy_family.py \ # from its dry run and spreads their files over parallel runners by recent # CI times, splitting a long group into chunks. A process still runs files # of one group only, in this order, so none holds more than its group here. -REST_PYTHON_GROUPS := cached-arrays heavy-a heavy-b core-heavy core policy-parameters policy remaining-heavy remaining +REST_PYTHON_GROUPS := cached-arrays heavy-a heavy-b core-heavy social-security-inputs core policy-parameters policy remaining-heavy remaining REST_CACHED_ARRAY_TESTS := $(TESTS)/test_formulas_do_not_write_into_cached_arrays.py REST_HEAVY_A_TESTS := $(TESTS)/test_md_poverty_line_credit_invariants.py \ $(TESTS)/core/test_ald_determinism.py \ @@ -69,6 +69,11 @@ REST_HEAVY_B_TESTS := $(TESTS)/core/test_marginal_tax_rate_coverage.py \ REST_CORE_HEAVY_TESTS := $(TESTS)/core/test_ssi_resource_deeming_invariants.py \ $(TESTS)/test_federal_credit_limit_order.py \ $(TESTS)/test_pinned_branch_credit_dependents.py +# Synthetic household input/branch checks load no national dataset. Keep their +# retained reform systems in a separate process rather than growing core's heap. +# Local 4 GiB-capped run peaked at 2.385 GiB before its 1800s loaded-host cap; +# CI run 40318424511 passed all 126 cases in 312.62s (6:20 process wall). +REST_SOCIAL_SECURITY_INPUT_TESTS := $(TESTS)/core/test_social_security_inputs.py # These three parameter modules construct full reform trees; one also loads # a default Microsimulation only to read parameters. Release their heap before # the remaining policy tests, rather than stacking both in the 13.9 GB group. @@ -83,7 +88,7 @@ REST_REMAINING_HEAVY_TESTS := $(TESTS)/test_niit_section_911_magi.py \ $(TESTS)/test_capital_gain_distributions_loss_limit.py \ $(TESTS)/test_schedule_d_filer_net_capital_gain.py REST_LISTED_TESTS := $(REST_CACHED_ARRAY_TESTS) $(REST_HEAVY_A_TESTS) $(REST_HEAVY_B_TESTS) \ - $(REST_CORE_HEAVY_TESTS) $(REST_POLICY_PARAMETER_TESTS) $(REST_REMAINING_HEAVY_TESTS) + $(REST_CORE_HEAVY_TESTS) $(REST_SOCIAL_SECURITY_INPUT_TESTS) $(REST_POLICY_PARAMETER_TESTS) $(REST_REMAINING_HEAVY_TESTS) REST_IGNORES := $(addprefix --ignore=,$(TESTS)/policy/contrib $(TESTS)/microsimulation $(REST_SPM_TESTS)) # $(call rest_pytest,,,) is the pytest # command for one group. With REST_REPORT_DIR set, each group writes its own @@ -310,6 +315,8 @@ test-other-python-rest-heavy-b: $(call rest_pytest,heavy-b,$(REST_HEAVY_B_TESTS)) test-other-python-rest-core-heavy: $(call rest_pytest,core-heavy,$(REST_CORE_HEAVY_TESTS)) +test-other-python-rest-social-security-inputs: + $(call rest_pytest,social-security-inputs,$(REST_SOCIAL_SECURITY_INPUT_TESTS)) test-other-python-rest-core: $(call rest_pytest,core,$(TESTS)/core,$(REST_LISTED_TESTS)) test-other-python-rest-policy-parameters: @@ -320,7 +327,7 @@ test-other-python-rest-remaining-heavy: $(call rest_pytest,remaining-heavy,$(REST_REMAINING_HEAVY_TESTS)) test-other-python-rest-remaining: $(call rest_pytest,remaining,$(TESTS)/,$(REST_LISTED_TESTS) $(TESTS)/core $(TESTS)/policy) -MICROSIMULATION_GROUPS := base social-security-inputs social-security-current-law +MICROSIMULATION_GROUPS := base social-security-current-law test-microsimulation: @failed=""; for group in $(MICROSIMULATION_GROUPS); do \ $(MAKE) --no-print-directory test-microsimulation-$$group || failed="$$failed $$group"; \ @@ -329,11 +336,11 @@ test-microsimulation: test-microsimulation-base: # Compatibility checks copy policy for differential/branch isolation. Keep # their heaps separate on this existing runner (local differential peak 7 GB). - pytest $(TESTS)/microsimulation --maxfail=0 --ignore=$(TESTS)/microsimulation/test_social_security_inputs.py --ignore=$(TESTS)/microsimulation/test_social_security_current_law.py -test-microsimulation-social-security-inputs: - pytest $(TESTS)/microsimulation/test_social_security_inputs.py --maxfail=0 $(if $(findstring --junitxml,$(PYTEST_ADDOPTS)),--junitxml=simulation-social-security-inputs.xml) + pytest $(TESTS)/microsimulation --maxfail=0 --ignore=$(TESTS)/microsimulation/test_social_security_current_law.py test-microsimulation-social-security-current-law: - pytest $(TESTS)/microsimulation/test_social_security_current_law.py --maxfail=0 $(if $(findstring --junitxml,$(PYTEST_ADDOPTS)),--junitxml=simulation-social-security-current-law.xml) + # Only the two probes need a country system. Avoid importing another through + # the empty parent conftest in the pytest process that compares their JSON. + pytest $(TESTS)/microsimulation/test_social_security_current_law.py --confcutdir=$(TESTS)/microsimulation --maxfail=0 $(if $(findstring --junitxml,$(PYTEST_ADDOPTS)),--junitxml=simulation-social-security-current-law.xml) # Local convenience: run both halves back to back. test-other: test-other-python test-microsimulation test-policy-contrib-python: diff --git a/policyengine_us/tests/code_health/test_rest_python_groups.py b/policyengine_us/tests/code_health/test_rest_python_groups.py index 5433b2278ef..47a455b216b 100644 --- a/policyengine_us/tests/code_health/test_rest_python_groups.py +++ b/policyengine_us/tests/code_health/test_rest_python_groups.py @@ -938,41 +938,67 @@ def stub_environment(tmp_path, outcomes): @pytest.mark.parametrize("failed_index", [None, 0, 1, 2]) -def test_microsimulation_processes_keep_reports_and_run_after_failures( - tmp_path, failed_index +def test_social_security_processes_keep_reports_and_run_after_failures( + tmp_path, groups, failed_index ): + """Keep three isolated processes after moving synthetic inputs to Rest.""" report_option = "--junitxml=simulation-microsimulation.xml" - commands = dry_run("test-microsimulation", f"PYTEST_ADDOPTS={report_option}") + microsimulation_commands = dry_run( + "test-microsimulation", f"PYTEST_ADDOPTS={report_option}" + ) + assert len(microsimulation_commands) == 2 + input_target = "test-other-python-rest-social-security-inputs" + (input_command,) = dry_run(input_target) + input_group = groups["social-security-inputs"] + assert input_command == ["pytest", *input_group["args"]] + commands = [*microsimulation_commands, input_command] assert len(commands) == 3 reports = [] - groups = {} - for index, command in enumerate(commands): + process_files = {} + for index, command in enumerate(microsimulation_commands): args = [report_option, *command[1:]] reports.append([a for a in args if a.startswith("--junitxml=")][-1]) - groups[str(index)] = collected_files( + process_files[str(index)] = collected_files( [a for a in command[1:] if not a.startswith("--junitxml=")] ) - reference = collected_files([str(TESTS / "microsimulation")]) - assert partition_problems(reference, groups) == [] - assert len(set(reports)) == 3 + input_file = "policyengine_us/tests/core/test_social_security_inputs.py" + process_files["inputs"] = collected_files(input_command[1:]) + assert process_files["inputs"] == [input_file] + reference = collected_files([str(TESTS / "microsimulation"), input_file]) + assert partition_problems(reference, process_files) == [] assert all( fnmatch(report.split("=", 1)[1], "simulation-*.xml") for report in reports ) + (input_report,) = [ + report for report in input_group["reports"] if report.endswith(".xml") + ] + assert input_report == f"{REPORT_DIR}/rest-python-social-security-inputs.xml" + reports.append(f"--junitxml={input_report}") + assert len(set(reports)) == 3 outcomes = ["0"] * 3 if failed_index is not None: outcomes[failed_index] = "1" env, calls = stub_environment(tmp_path, outcomes) env["PYTEST_ADDOPTS"] = report_option - result = subprocess.run( + microsimulation_result = subprocess.run( ["make", "--no-print-directory", "test-microsimulation"], cwd=REPO, env=env, capture_output=True, text=True, ) + env.pop("PYTEST_ADDOPTS") + input_result = subprocess.run( + ["make", "--no-print-directory", input_target], + cwd=REPO, + env=env, + capture_output=True, + text=True, + ) assert calls.read_text().splitlines() == ["|".join(c[1:]) for c in commands] - assert (result.returncode != 0) == (failed_index is not None) + assert (microsimulation_result.returncode != 0) == (failed_index in (0, 1)) + assert (input_result.returncode != 0) == (failed_index == 2) def scenario_outcomes(scenario, count): diff --git a/policyengine_us/tests/microsimulation/test_social_security_inputs.py b/policyengine_us/tests/core/test_social_security_inputs.py similarity index 100% rename from policyengine_us/tests/microsimulation/test_social_security_inputs.py rename to policyengine_us/tests/core/test_social_security_inputs.py diff --git a/policyengine_us/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py index 138254d09e3..f07988fd6b2 100644 --- a/policyengine_us/tests/microsimulation/test_social_security_current_law.py +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -994,7 +994,9 @@ def differential(tmp_path_factory): env=environment, capture_output=True, text=True, - timeout=420, + # Each 74-route probe took about 290s in CI run 40318424511. + # Allow slower hosts while retaining a bounded hang guard. + timeout=900, check=False, ) assert completed.returncode == 0, ( From 13216f9be93e6ce844692526bb26f7a74e2a3981 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 17:52:15 -0700 Subject: [PATCH 56/76] Fix refresh of Social Security formula returns without holder storage --- changelog.d/ss-benefit-levers.added.md | 2 +- .../core/test_social_security_real_core.py | 223 +++++++++++++++++- .../tools/social_security_inputs.py | 91 ++++++- 3 files changed, 306 insertions(+), 10 deletions(-) diff --git a/changelog.d/ss-benefit-levers.added.md b/changelog.d/ss-benefit-levers.added.md index f04f3a13901..2a7fe02d748 100644 --- a/changelog.d/ss-benefit-levers.added.md +++ b/changelog.d/ss-benefit-levers.added.md @@ -1 +1 @@ -Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs remain independent sources, exposed through a read-time reported view without automatic twin arrays or provenance stamps. Initial loading routes sources so same-year benefit factors apply; later public writes remain final overrides. Explicit reported inputs retain priority, while exact compound totals keep their period semantics. Compound helpers retain initial-input precedence. Source refreshes use Core's dependency sequence cutoff, preserving raw caller values, supplied inputs and earlier formula results. Builders, partial loading failures, copies, cache replay, extraction, dump restoration and subsampling preserve successful source inputs. A state machine searches combinations of write and copy routes. +Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs remain independent sources, exposed through a read-time reported view without automatic twin arrays or provenance stamps. Initial loading routes sources so same-year benefit factors apply; later public writes remain final overrides. Explicit reported inputs retain priority, while exact compound totals keep their period semantics. Compound helpers retain initial-input precedence. Source refreshes use Core's dependency sequence cutoff, including formula returns kept only in the fast cache, preserving raw caller values, supplied inputs and earlier formula results. Builders, partial loading failures, copies, cache replay, extraction, dump restoration and subsampling preserve successful source inputs. A state machine searches combinations of write and copy routes. diff --git a/policyengine_us/tests/core/test_social_security_real_core.py b/policyengine_us/tests/core/test_social_security_real_core.py index 9825d821d00..4f51df43844 100644 --- a/policyengine_us/tests/core/test_social_security_real_core.py +++ b/policyengine_us/tests/core/test_social_security_real_core.py @@ -15,6 +15,7 @@ import numpy as np import pytest from policyengine_core.entities import build_entity +from policyengine_core.experimental import MemoryConfig from policyengine_core.holders import Holder, helpers from policyengine_core.holders.helpers import ( set_input_dispatch_by_period, @@ -26,6 +27,7 @@ from policyengine_core.taxbenefitsystems import TaxBenefitSystem from policyengine_core.tools import simulation_dumper from policyengine_core.variables import Variable +from policyengine_core.warnings import MemoryConfigWarning COUNTRY_ROOT = Path(__file__).resolve().parents[2] @@ -166,6 +168,54 @@ def fresh(inputs=None, *, enabled=True, factor=1): ) +def fresh_without_holder_cache(mode, names, inputs=None): + """Use each supported opt-out during Core's Holder initialization.""" + system = make_system() + simulation = Simulation( + tax_benefit_system=system, populations=system.instantiate_entities() + ) + if mode == "blacklist": + system.cache_blacklist = list(names) + simulation.opt_out_cache = True + else: + assert mode == "variables_to_drop" + with pytest.warns(MemoryConfigWarning): + simulation.memory_config = MemoryConfig( + 1, + # The tests exercise dropping, without introducing disk I/O. + priority_variables=list(system.variables), + variables_to_drop=names, + ) + # Simulation.__init__ creates empty holders while enumerating its input + # variables. Replace those populations so later Holder initialization + # reads the actual memory_config rather than its constructor default. + simulation.build_from_populations(system.instantiate_entities()) + values = {"independent_input": {"2023": [70.0], "2024": [70.0]}} + values.update(inputs or {}) + return SimulationBuilder().build_from_variables( + system, values, simulation=simulation + ) + + +def rebuild_with_source(simulation, value): + from policyengine_core.data import Dataset + + class TinyDataset: + data_format = Dataset.TIME_PERIOD_ARRAYS + time_period = "2024" + + def load(self): + return { + "person_id": {"2024": np.array([1])}, + "household_id": {"2024": np.array([1])}, + "person_household_id": {"2024": np.array([1])}, + PUBLIC: {"2024": np.array([value])}, + } + + simulation.dataset = TinyDataset() + simulation.build_from_dataset() + + def scalar(simulation, name, requested): return float(simulation.calculate(name, requested)[0]) @@ -422,8 +472,175 @@ def load(self): assert scalar(simulation, PUBLIC, "2023") == 1 +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +@pytest.mark.parametrize("route", ("dataset", "set_input", "branch_set")) +@pytest.mark.parametrize("initial", (0, 12_000), ids=("default_zero", "initial_source")) +def test_source_write_refreshes_fast_only_formula_returns(mode, route, initial): + names = [PUBLIC, PUBLIC + "_base"] + if not initial: + # The review's zero-source variant also opts the reported default out. + names.append(PUBLIC + "_reported") + inputs = {PUBLIC: {"2024": [initial]}} if initial else {} + simulation = fresh_without_holder_cache(mode, names, inputs) + if route == "branch_set": + simulation = simulation.get_branch("child") + assert scalar(simulation, PUBLIC, "2024") == initial + assert earliest_source_dependency(simulation) is not None + assert stored(simulation, PUBLIC, "2024") is None + assert stored(simulation, PUBLIC + "_base", "2024") is None + key = (PUBLIC, period("2024")) + assert float(simulation._fast_cache[key][0]) == initial + epoch = simulation._input_epoch + if route == "dataset": + rebuild_with_source(simulation, 24_000.0) + elif route == "set_input": + simulation.set_input(PUBLIC, "2024", [24_000.0]) + else: + # Exercise the adapter itself: Core adds dependency cleanup only on + # Holder.set_input, not on this supported lower-level branch route. + simulation.get_holder(PUBLIC)._set( + period("2024"), + np.array([24_000.0]), + simulation.branch_name, + is_input=True, + ) + assert key not in simulation._fast_cache + assert simulation._input_epoch > epoch + assert scalar(simulation, PUBLIC + "_reported", "2024") == 24_000 + assert scalar(simulation, PUBLIC, "2024") == 24_000 + + +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +def test_fast_only_formula_before_since_survives_without_a_drop_witness(mode): + simulation = fresh_without_holder_cache(mode, ["independent_formula"]) + assert scalar(simulation, "independent_formula", "2023") == 77 + key = ("independent_formula", period("2023")) + before = simulation._fast_cache[key] + assert stored(simulation, "independent_formula", "2023") is None + number = simulation._get_store_history().earliest_dependency(*key) + simulation.set_input(PUBLIC, "2024", [12_000.0]) + assert scalar(simulation, PUBLIC + "_reported", "2024") == 12_000 + assert number < earliest_source_dependency(simulation) + epoch = simulation._input_epoch + simulation.set_input(PUBLIC, "2024", [24_000.0]) + # A pre-cutoff fast-only return cannot itself justify a global clear. + assert simulation._fast_cache[key] is before + assert simulation._input_epoch == epoch + assert scalar(simulation, "independent_formula", "2023") == 77 + + +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +def test_fast_only_unrelated_formula_after_since_is_intentionally_dropped(mode): + simulation = fresh_without_holder_cache( + mode, ["independent_formula"], {PUBLIC: {"2024": [12_000.0]}} + ) + assert scalar(simulation, PUBLIC + "_reported", "2024") == 12_000 + assert scalar(simulation, "independent_formula", "2024") == 77 + key = ("independent_formula", period("2024")) + assert stored(simulation, *key) is None + assert simulation._get_store_history().earliest_dependency( + *key + ) >= earliest_source_dependency(simulation) + epoch = simulation._input_epoch + simulation.set_input(PUBLIC, "2024", [24_000.0]) + # Intended over-approximation: refresh-scope ruling rule 5 includes an + # unrelated formula return whose own store number is after the cutoff. + assert key not in simulation._fast_cache + assert simulation._input_epoch > epoch + assert scalar(simulation, "independent_formula", "2024") == 77 + + +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +def test_recalculated_fast_only_return_uses_its_own_store_number(mode, monkeypatch): + stores = [] + original_record_store = Holder._record_store + + def record_store(holder, requested, number): + if holder.variable.name == "independent_formula": + stores.append(number) + return original_record_store(holder, requested, number) + + monkeypatch.setattr(Holder, "_record_store", record_store) + simulation = fresh_without_holder_cache(mode, ["independent_formula"]) + assert scalar(simulation, "independent_formula", "2023") == 77 + key = ("independent_formula", period("2023")) + first = simulation._get_store_history().earliest_dependency(*key) + assert first == stores[-1] + simulation.set_input(PUBLIC, "2024", [12_000.0]) + assert scalar(simulation, PUBLIC + "_reported", "2024") == 12_000 + since = earliest_source_dependency(simulation) + assert first < since + # Recompute the same key after the source cutoff. StoreHistory keeps its + # first number, whereas the memo must be dated by this later calculation. + simulation._fast_cache.pop(key) + assert scalar(simulation, "independent_formula", "2023") == 77 + assert stores[-1] >= since + assert simulation._get_store_history().earliest_dependency(*key) == first + epoch = simulation._input_epoch + simulation.set_input(PUBLIC, "2024", [24_000.0]) + assert key not in simulation._fast_cache + assert simulation._input_epoch > epoch + + +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +@pytest.mark.parametrize("branch", (False, True), ids=("main", "branch")) +@pytest.mark.parametrize("registered", (False, True), ids=("raw_seed", "input")) +def test_not_kept_formula_name_preserves_physical_caller_values( + mode, branch, registered +): + names = [PUBLIC, PUBLIC + "_base", "independent_formula"] + simulation = fresh_without_holder_cache(mode, names, {PUBLIC: {"2024": [12_000.0]}}) + if branch: + simulation = simulation.get_branch("child") + assert scalar(simulation, PUBLIC, "2024") == 12_000 + holder = simulation.get_holder("independent_formula") + if registered: + simulation.set_input("independent_formula", "2025", [99.0]) + else: + # _set can supply a raw seed even where put_in_cache is opted out. + holder._set( + period("2025"), + np.array([99.0]), + simulation.branch_name, + derived=False, + is_input=False, + ) + before = stored(simulation, "independent_formula", "2025") + assert sequence_number( + simulation, "independent_formula", "2025" + ) >= earliest_source_dependency(simulation) + assert not holder.is_derived(period("2025"), simulation.branch_name) + epoch = simulation._input_epoch + simulation.get_holder(PUBLIC)._set( + period("2024"), + np.array([24_000.0]), + simulation.branch_name, + is_input=True, + ) + assert simulation._input_epoch > epoch + assert stored(simulation, "independent_formula", "2025") is before + assert holder._is_input(period("2025"), simulation.branch_name) is registered + assert scalar(simulation, "independent_formula", "2025") == 99 + assert scalar(simulation, PUBLIC, "2024") == 24_000 + + +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +def test_unread_source_load_preserves_fast_only_returns_without_a_cutoff(mode): + names = [PUBLIC, PUBLIC + "_base", PUBLIC + "_reported", "independent_formula"] + simulation = fresh_without_holder_cache(mode, names) + assert scalar(simulation, "independent_formula", "2023") == 77 + key = ("independent_formula", period("2023")) + before = simulation._fast_cache[key] + assert earliest_source_dependency(simulation) is None + epoch = simulation._input_epoch + rebuild_with_source(simulation, 24_000.0) + assert simulation._input_epoch == epoch + assert simulation._fast_cache[key] is before + assert scalar(simulation, PUBLIC, "2024") == 24_000 + + @pytest.mark.parametrize("initial", (False, True), ids=("raw_source", "initial_source")) -def test_public_write_evicts_exact_reported_memo_without_dropping_other_caches(initial): +def test_exact_reported_read_bypasses_native_memo_and_keeps_other_caches(initial): simulation = fresh() assert scalar(simulation, "independent_formula", "2024") == 77 independent_key = ("independent_formula", period("2024")) @@ -439,6 +656,10 @@ def test_public_write_evicts_exact_reported_memo_without_dropping_other_caches(i np.array([12_000.0]), period("2024"), derived=False ) assert scalar(simulation, PUBLIC + "_reported", "2024") == 12_000 + # Native Core returns this exact source directly from get_array; unlike + # the synthetic backend, neither supported Core version inserts a memo. + reported_key = (PUBLIC + "_reported", period("2024")) + assert reported_key not in simulation._fast_cache # This exact read returns the source view without storing a derived # reported array. There is no post-cutoff formula array to purge. assert ( diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index e2255d01e32..af2a96f07bb 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -18,7 +18,7 @@ from dataclasses import dataclass from functools import wraps from pathlib import Path -from weakref import ref +from weakref import WeakKeyDictionary, ref from policyengine_core.data_storage.store_history import next_sequence_number from policyengine_core.periods import period as to_period @@ -48,6 +48,11 @@ _subsampling_social_security_sources = ContextVar( "subsampling_ss_sources", default=None ) +# Clone copies unknown attributes shallowly. Keep return metadata outside the +# simulation so independent branches never share a mutable metadata dictionary. +# Preserve it when importlib.reload leaves the installed Core hooks in place. +_fast_only_returns = globals().get("_fast_only_returns", WeakKeyDictionary()) +_cache_store_number = ContextVar("ss_cache_store_number", default=None) _PHYSICAL = {} @@ -528,6 +533,12 @@ def has_source(holder, period, branches): continue branch = holder._branch_storing(period, simulation.branch_name) if branch is None: + # StoreHistory remembers the first store, not the store which + # produced this return. Skipped-storage formula returns need their + # own number, captured by _cache_result/_record_store below. + entry = _fast_only_returns.get(simulation, {}).get((name, period)) + if entry is not None and entry[1]() is array: + entries[key] = entry continue if holder._memory_storage.has(period, branch): storage = holder._memory_storage @@ -562,6 +573,8 @@ def _purge_derived_arrays(simulation, public, period): if not dependencies: return 0 since = min(dependencies) + # Collect before physical deletion changes the holders' slot metadata. + cached_returns = _derived_cached_returns(simulation) dropped = 0 for population in simulation.populations.values(): for holder in population._holders.values(): @@ -585,15 +598,18 @@ def _purge_derived_arrays(simulation, public, period): ): _delete_exact(holder, period, branch, tier=tier) dropped += 1 - for owner, entries in _loading_social_security_cache_entries.get(): + for owner, entries in ( + (simulation, cached_returns), + *_loading_social_security_cache_entries.get(), + ): if owner is not simulation: continue for key, (number, reference) in tuple(entries.items()): if number >= since: array = reference() if array is not None and simulation._fast_cache.get(key) is array: - # This numbered formula array remains in the fast cache - # after Core detached its original holder during rebuild. + # A numbered formula return can outlive its holder, or + # Core can skip holder storage while memoizing the return. simulation._fast_cache.pop(key) dropped += 1 entries.pop(key) @@ -644,9 +660,68 @@ def install_social_security_input_compatibility(): recorded_periods=CoreSimulation._get_set_input_periods, recorded_array=CoreSimulation._get_recorded_input_array, calculate=CoreSimulation._calculate, + cache_result=CoreSimulation._cache_result, + record_store=Holder._record_store, stored_input=helpers.get_stored_input, ) + @wraps(_PHYSICAL["record_store"]) + def record_store(holder, period, sequence_number): + result = _PHYSICAL["record_store"](holder, period, sequence_number) + capture = _cache_store_number.get() + if ( + capture is not None + and capture[0]() is holder + and capture[1] == holder._storage_period(period) + ): + capture[2].append(sequence_number) + return result + + Holder._record_store = record_store + + @wraps(_PHYSICAL["cache_result"]) + def cache_result(simulation, holder, array, period, input_state): + if not _enabled(simulation) or not ( + holder._do_not_store + or ( + simulation.opt_out_cache + and simulation.tax_benefit_system.cache_blacklist + and holder.variable.name + in simulation.tax_benefit_system.cache_blacklist + ) + ): + return _PHYSICAL["cache_result"]( + simulation, holder, array, period, input_state + ) + stored_period = holder._storage_period(period) + numbers = [] + token = _cache_store_number.set((ref(holder), stored_period, numbers)) + try: + result = _PHYSICAL["cache_result"]( + simulation, holder, array, period, input_state + ) + finally: + _cache_store_number.reset(token) + # _cache_result can instead return a newly supplied input. With no + # formula store there is no number and no derived-return witness. + if ( + numbers + and holder._branch_storing(stored_period, simulation.branch_name) is None + ): + entries = _fast_only_returns.setdefault(simulation, {}) + key = (holder.variable.name, stored_period) + number = numbers[-1] + + def forget(reference): + current = entries.get(key) + if current is not None and current[1] is reference: + entries.pop(key, None) + + entries[key] = (number, ref(result, forget)) + return result + + CoreSimulation._cache_result = cache_result + @wraps(_PHYSICAL["stored_input"]) def get_stored_input(holder, period, branch_name): value = _PHYSICAL["stored_input"](holder, period, branch_name) @@ -821,10 +896,10 @@ def store( if name in SOCIAL_SECURITY_COMPONENTS: reported_holder = _holder(simulation, f"{name}_reported") if reported_holder is not None: - # An exact reported read can memoize the public input - # itself, without storing any derived array. Evict that - # changed source view exactly as Core evicts the written - # holder's own memo; unrelated cached returns stay intact. + # Defensive for calculation backends that memoize exact + # source reads. Native Core 3.33.2/3.33.12 returns these + # before memo insertion. Evict only this changed view, + # as Core does for the written holder's own memo. reported_holder._evict_fast_cache( holder._storage_period(period), effective_branch ) From fb4d23883191e9d2dac8fa9cdc90f99dfc7d2dfb Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 17:53:49 -0700 Subject: [PATCH 57/76] Generate skipped-storage cache policies in Social Security input machine --- .../core/_social_security_state_backend.py | 24 +++++++-- ...est_social_security_input_state_machine.py | 54 ++++++++++++++++--- 2 files changed, 67 insertions(+), 11 deletions(-) diff --git a/policyengine_us/tests/core/_social_security_state_backend.py b/policyengine_us/tests/core/_social_security_state_backend.py index 706f469db6a..c60b28809e5 100644 --- a/policyengine_us/tests/core/_social_security_state_backend.py +++ b/policyengine_us/tests/core/_social_security_state_backend.py @@ -50,7 +50,7 @@ from importlib.metadata import distribution from pathlib import Path from unittest.mock import patch -from weakref import ref +from weakref import WeakKeyDictionary, ref import numpy as np from policyengine_core import commons, periods @@ -163,8 +163,10 @@ def get(self, period): class Backend: """Install the unchanged hooks around unchanged Core routes in isolation.""" - def __init__(self): + def __init__(self, *, cache_blacklist=(), variables_to_drop=()): self.stack = ExitStack() + self.cache_blacklist = tuple(cache_blacklist) + self.variables_to_drop = tuple(variables_to_drop) try: self._setup() except BaseException: @@ -172,6 +174,8 @@ def __init__(self): raise def _setup(self): + cache_blacklist = self.cache_blacklist + variables_to_drop = self.variables_to_drop core = Path(distribution("policyengine-core").locate_file("policyengine_core")) calculation_frames = core_simulation._calculation_frames formula_simulation = core_simulation._formula_simulation @@ -198,7 +202,9 @@ def __init__(self, population, name): self._memory_storage = InMemoryStorage(is_eternal=False) self._disk_storage = None self._on_disk_storable = False - self._do_not_store = False + self._do_not_store = name in ( + self.simulation.memory_config.variables_to_drop + ) def _storage_period(self, period): return as_period(period) @@ -255,11 +261,11 @@ class System: _social_security_input_compatibility = True person_entity = person group_entities = [household] - cache_blacklist = [] data_modified = False auto_carry_over_input_variables = True def __init__(self): + self.cache_blacklist = list(cache_blacklist) self.variables = { name: types.SimpleNamespace( name=name, @@ -292,6 +298,7 @@ def __init__(self): "raw_seed", "early_derived", "late_derived", + "unrelated_formula", *COMPONENTS, *(f"{name}_reported" for name in COMPONENTS), *(f"{name}_base" for name in COMPONENTS), @@ -342,7 +349,11 @@ def __init__(self, tax_benefit_system=None, populations=None, **kwargs): self.invalidated_caches = set() self._fast_cache = {} self.start_instant = None - self.opt_out_cache = False + self.opt_out_cache = bool(cache_blacklist) + self.memory_config = types.SimpleNamespace( + variables_to_drop=set(variables_to_drop), + priority_variables=set(self.tax_benefit_system.variables), + ) self.trace = False self.fail_after_store = False self.build_from_populations( @@ -396,6 +407,8 @@ def _check_for_cycle(self, *args): pass def _run_formula(self, variable, population, period): + if variable.name == "unrelated_formula": + return np.array([77.0]) if variable.name in COMPONENTS: return self.calculate(f"{variable.name}_base", period) * 2 if variable.name.endswith("_base"): @@ -727,6 +740,7 @@ def module(name, **values): "_social_security_state_hooks", json=json, ContextVar=ContextVar, + WeakKeyDictionary=WeakKeyDictionary, ref=ref, dataclass=dataclass, wraps=wraps, diff --git a/policyengine_us/tests/core/test_social_security_input_state_machine.py b/policyengine_us/tests/core/test_social_security_input_state_machine.py index 67f3e5dc56f..eca7c39ce12 100644 --- a/policyengine_us/tests/core/test_social_security_input_state_machine.py +++ b/policyengine_us/tests/core/test_social_security_input_state_machine.py @@ -24,7 +24,13 @@ from pathlib import Path from hypothesis import HealthCheck, settings, strategies as st -from hypothesis.stateful import RuleBasedStateMachine, invariant, precondition, rule +from hypothesis.stateful import ( + RuleBasedStateMachine, + initialize, + invariant, + precondition, + rule, +) spec = importlib.util.spec_from_file_location( @@ -42,6 +48,16 @@ PERIOD = st.sampled_from(("2023", "2024", "2025", "year:2023:2", "year:2024:2")) AMOUNT = st.integers(min_value=0, max_value=48_000) TARGET = st.integers(min_value=0, max_value=31) +CACHE_NAMES = ( + *COMPONENTS, + *(f"{name}_base" for name in COMPONENTS), + *(f"{name}_reported" for name in COMPONENTS), + "unrelated_formula", +) +CACHE_CONFIGURATION = st.tuples( + st.sampled_from(("blacklist", "variables_to_drop")), + st.sets(st.sampled_from(CACHE_NAMES), min_size=1), +) @dataclass @@ -233,17 +249,28 @@ class SocialSecurityInputsMachine(RuleBasedStateMachine): def __init__(self): super().__init__() - self.backend = Backend() - self.states = [(self.backend.DelegatingSimulation(), Model())] + self.backend = None + self.states = [] self.temp = tempfile.TemporaryDirectory(prefix="ss-input-machine-") self.dump_count = 0 self.branch_initial_rewrite_done = False self.preservation_checks = [] self.helper_checks = [] + @initialize(cache_configuration=CACHE_CONFIGURATION) + def configure_cache(self, cache_configuration): + """Search Core's two policies that number returns without storing arrays.""" + mode, names = cache_configuration + self.backend = Backend( + cache_blacklist=names if mode == "blacklist" else (), + variables_to_drop=names if mode == "variables_to_drop" else (), + ) + self.states = [(self.backend.DelegatingSimulation(), Model())] + def teardown(self): self.temp.cleanup() - self.backend.close() + if self.backend is not None: + self.backend.close() def state(self, slot): return self.states[slot % len(self.states)] @@ -596,10 +623,21 @@ def low_level_store( def nonderived_cache(self, slot, component, reported, period, amount, data=None): simulation, model = self.state(slot) component, period = self.write_key(model, component, period, data) - simulation.get_holder(self.name(component, reported)).put_in_cache( + holder = simulation.get_holder(self.name(component, reported)) + holder.put_in_cache( [amount], as_period(period), simulation.branch_name, derived=False ) - model.write(component, reported, period, amount, False, "put_in_cache") + # Core's put_in_cache ignores nonderived seeds too under these + # policies. A raw _set still stores and is covered by low_level_store. + if not ( + holder._do_not_store + or ( + simulation.opt_out_cache + and holder.variable.name + in simulation.tax_benefit_system.cache_blacklist + ) + ): + model.write(component, reported, period, amount, False, "put_in_cache") @rule( slot=TARGET, @@ -869,6 +907,10 @@ def reported_values_and_provenance_follow_intent(self): for simulation, model in self.states: assert simulation._user_input_contexts == [] assert not hasattr(simulation, "_social_security_automatic_reported_inputs") + assert ( + float(simulation.calculate("unrelated_formula", as_period("2024"))[0]) + == 77 + ) for key in model.keys(): component, period = key expected, expected_automatic = model.reported(component, period) From 06092f17bf591040931f4408043e6883bc8d541d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 18:17:39 -0700 Subject: [PATCH 58/76] Cover reported-source writes with fast-only benefit caches --- .../core/test_social_security_real_core.py | 32 +++++++++++++++---- 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/policyengine_us/tests/core/test_social_security_real_core.py b/policyengine_us/tests/core/test_social_security_real_core.py index 4f51df43844..d1a652741b3 100644 --- a/policyengine_us/tests/core/test_social_security_real_core.py +++ b/policyengine_us/tests/core/test_social_security_real_core.py @@ -473,9 +473,25 @@ def load(self): @pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) -@pytest.mark.parametrize("route", ("dataset", "set_input", "branch_set")) +@pytest.mark.parametrize( + "route,source", + ( + ("dataset", PUBLIC), + ("set_input", PUBLIC), + ("set_input", PUBLIC + "_reported"), + ("branch_set", PUBLIC), + ("branch_set", PUBLIC + "_reported"), + ), + ids=( + "dataset_public", + "set_public", + "set_reported", + "branch_public", + "branch_reported", + ), +) @pytest.mark.parametrize("initial", (0, 12_000), ids=("default_zero", "initial_source")) -def test_source_write_refreshes_fast_only_formula_returns(mode, route, initial): +def test_source_write_refreshes_fast_only_formula_returns(mode, route, source, initial): names = [PUBLIC, PUBLIC + "_base"] if not initial: # The review's zero-source variant also opts the reported default out. @@ -494,20 +510,22 @@ def test_source_write_refreshes_fast_only_formula_returns(mode, route, initial): if route == "dataset": rebuild_with_source(simulation, 24_000.0) elif route == "set_input": - simulation.set_input(PUBLIC, "2024", [24_000.0]) + simulation.set_input(source, "2024", [24_000.0]) else: # Exercise the adapter itself: Core adds dependency cleanup only on # Holder.set_input, not on this supported lower-level branch route. - simulation.get_holder(PUBLIC)._set( + simulation.get_holder(source)._set( period("2024"), np.array([24_000.0]), simulation.branch_name, is_input=True, ) - assert key not in simulation._fast_cache + memo_survived_write = key in simulation._fast_cache + reported_after = scalar(simulation, PUBLIC + "_reported", "2024") + public_after = scalar(simulation, PUBLIC, "2024") + assert (reported_after, public_after) == (24_000, 24_000) + assert not memo_survived_write assert simulation._input_epoch > epoch - assert scalar(simulation, PUBLIC + "_reported", "2024") == 24_000 - assert scalar(simulation, PUBLIC, "2024") == 24_000 @pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) From 871ee0eeab1c8850b72e7242e021abbe70ed0488 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 18:40:51 -0700 Subject: [PATCH 59/76] Date fast returns at memo insertion and preserve cache serialization --- .../core/test_social_security_real_core.py | 64 ++++++++++ .../tools/social_security_inputs.py | 112 +++++++++++++++--- 2 files changed, 159 insertions(+), 17 deletions(-) diff --git a/policyengine_us/tests/core/test_social_security_real_core.py b/policyengine_us/tests/core/test_social_security_real_core.py index d1a652741b3..6c0d83361b8 100644 --- a/policyengine_us/tests/core/test_social_security_real_core.py +++ b/policyengine_us/tests/core/test_social_security_real_core.py @@ -8,6 +8,7 @@ from __future__ import annotations import importlib +import pickle import sys import types from pathlib import Path @@ -600,6 +601,69 @@ def record_store(holder, requested, number): assert simulation._input_epoch > epoch +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +@pytest.mark.parametrize("memo_before_since", (False, True), ids=("after", "before")) +@pytest.mark.parametrize("reuse_array", (False, True), ids=("new_array", "same_array")) +def test_nonmemoized_store_keeps_the_existing_fast_return_date( + mode, memo_before_since, reuse_array +): + inputs = {} if memo_before_since else {PUBLIC: {"2024": [12_000.0]}} + simulation = fresh_without_holder_cache(mode, ["independent_formula"], inputs) + if not memo_before_since: + assert scalar(simulation, PUBLIC + "_reported", "2024") == 12_000 + assert scalar(simulation, "independent_formula", "2023") == 77 + key = ("independent_formula", period("2023")) + memo = simulation._fast_cache[key] + number = simulation._get_store_history().earliest_dependency(*key) + if memo_before_since: + simulation.set_input(PUBLIC, "2024", [12_000.0]) + assert scalar(simulation, PUBLIC + "_reported", "2024") == 12_000 + since = earliest_source_dependency(simulation) + assert (number < since) is memo_before_since + holder = simulation.get_holder("independent_formula") + assert stored(simulation, *key) is None + + # A helper can store a result through _cache_result without inserting it + # into the fast cache. Even reusing the same ndarray must not change the + # earlier memo's date: its producing store is the original calculation. + candidate = memo if reuse_array else memo.copy() + assert ( + simulation._cache_result( + holder, candidate, key[1], simulation._calculation_start() + ) + is candidate + ) + assert simulation._fast_cache[key] is memo + assert stored(simulation, *key) is None + epoch = simulation._input_epoch + simulation.set_input(PUBLIC, "2024", [24_000.0]) + if memo_before_since: + assert simulation._fast_cache[key] is memo + assert simulation._input_epoch == epoch + else: + # Intended rule 5 drop: the original memo itself followed the cutoff. + assert key not in simulation._fast_cache + assert simulation._input_epoch > epoch + + +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +def test_pickled_fast_only_cache_retains_its_return_dates(mode): + simulation = fresh_without_holder_cache( + mode, ["independent_formula"], {PUBLIC: {"2024": [12_000.0]}} + ) + assert scalar(simulation, PUBLIC + "_reported", "2024") == 12_000 + assert scalar(simulation, "independent_formula", "2024") == 77 + key = ("independent_formula", period("2024")) + assert stored(simulation, *key) is None + simulation._fast_cache = pickle.loads(pickle.dumps(simulation._fast_cache)) + np.testing.assert_array_equal(simulation._fast_cache[key], [77]) + assert stored(simulation, *key) is None + epoch = simulation._input_epoch + simulation.set_input(PUBLIC, "2024", [24_000.0]) + assert key not in simulation._fast_cache + assert simulation._input_epoch > epoch + + @pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) @pytest.mark.parametrize("branch", (False, True), ids=("main", "branch")) @pytest.mark.parametrize("registered", (False, True), ids=("raw_seed", "input")) diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index af2a96f07bb..b5abeeb83e4 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -48,10 +48,10 @@ _subsampling_social_security_sources = ContextVar( "subsampling_ss_sources", default=None ) -# Clone copies unknown attributes shallowly. Keep return metadata outside the -# simulation so independent branches never share a mutable metadata dictionary. -# Preserve it when importlib.reload leaves the installed Core hooks in place. -_fast_only_returns = globals().get("_fast_only_returns", WeakKeyDictionary()) +# Pending formula returns contain only numbers and weak references. Clone +# resets Core's fast cache; weak simulation keys keep pending stores independent. +# Preserve them across reload while installed Core hooks remain in place. +_pending_fast_returns = globals().get("_pending_fast_returns", WeakKeyDictionary()) _cache_store_number = ContextVar("ss_cache_store_number", default=None) _PHYSICAL = {} @@ -480,6 +480,82 @@ def _select_subsample_sources(simulation, metadata, person_ids): } +def _remember_weak_return(entries, key, number, array): + def forget(reference): + current = entries.get(key) + if current is not None and current[1] is reference: + entries.pop(key, None) + + entries[key] = (number, ref(array, forget)) + + +def _restore_formula_fast_cache(values, numbers): + return _FormulaFastCache(values, numbers) + + +class _FormulaFastCache(dict): + """Core's memo values with the actual store number of each insertion. + + A helper may number a return without memoizing it, even reusing an array + already in the memo. Only bracket assignment commits its new date. Values + have the same ownership as Core's dictionary; extra state holds numbers + and a weak owner only. Native clones/drops reset this cache to an empty dict. + """ + + def __init__(self, values=(), numbers=None): + super().__init__(values) + self._store_numbers = dict(numbers or {}) + self._owner = None + + def __setitem__(self, key, array): + super().__setitem__(key, array) + self._store_numbers.pop(key, None) + simulation = None if self._owner is None else self._owner() + if simulation is None: + return + name, requested = key + requested = None if requested is None else to_period(requested) + pending = _pending_fast_returns.get(simulation, {}) + entry = pending.pop((name, requested, id(array)), None) + if entry is not None and entry[1]() is array: + self._store_numbers[key] = entry[0] + + def __delitem__(self, key): + super().__delitem__(key) + self._store_numbers.pop(key, None) + + def pop(self, key, *default): + self._store_numbers.pop(key, None) + return super().pop(key, *default) + + def popitem(self): + key, array = super().popitem() + self._store_numbers.pop(key, None) + return key, array + + def clear(self): + super().clear() + self._store_numbers.clear() + + def update(self, *args, **kwargs): + for key, array in dict(*args, **kwargs).items(): + self[key] = array + + def setdefault(self, key, default=None): + if key not in self: + self[key] = default + return self[key] + + def __ior__(self, other): + self.update(other) + return self + + def __reduce__(self): + # Weak ownership/pending state is ephemeral. Preserve memo dates with + # the values, so pickle remains usable and a restored memo can refresh. + return _restore_formula_fast_cache, (dict(self), self._store_numbers) + + def _derived_cached_returns(simulation): """Remember numbered formula returns without retaining an old dataset. @@ -536,9 +612,9 @@ def has_source(holder, period, branches): # StoreHistory remembers the first store, not the store which # produced this return. Skipped-storage formula returns need their # own number, captured by _cache_result/_record_store below. - entry = _fast_only_returns.get(simulation, {}).get((name, period)) - if entry is not None and entry[1]() is array: - entries[key] = entry + number = getattr(fast_cache, "_store_numbers", {}).get(key) + if number is not None: + entries[key] = (number, ref(array)) continue if holder._memory_storage.has(period, branch): storage = holder._memory_storage @@ -708,16 +784,18 @@ def cache_result(simulation, holder, array, period, input_state): numbers and holder._branch_storing(stored_period, simulation.branch_name) is None ): - entries = _fast_only_returns.setdefault(simulation, {}) - key = (holder.variable.name, stored_period) - number = numbers[-1] - - def forget(reference): - current = entries.get(key) - if current is not None and current[1] is reference: - entries.pop(key, None) - - entries[key] = (number, ref(result, forget)) + cache = simulation._fast_cache + if not isinstance(cache, _FormulaFastCache): + cache = _FormulaFastCache(cache, getattr(cache, "_store_numbers", {})) + simulation._fast_cache = cache + cache._owner = ref(simulation) + requested = None if period is None else to_period(period) + _remember_weak_return( + _pending_fast_returns.setdefault(simulation, {}), + (holder.variable.name, requested, id(result)), + numbers[-1], + result, + ) return result CoreSimulation._cache_result = cache_result From d55a67302f6ac62f5e1ece90518fe23c9281773f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 18:45:22 -0700 Subject: [PATCH 60/76] Keep restored memo dates ahead of subsequent Core stores --- .../core/_social_security_state_backend.py | 2 ++ .../core/test_social_security_real_core.py | 32 +++++++++++++++++-- .../tools/social_security_inputs.py | 7 +++- 3 files changed, 38 insertions(+), 3 deletions(-) diff --git a/policyengine_us/tests/core/_social_security_state_backend.py b/policyengine_us/tests/core/_social_security_state_backend.py index c60b28809e5..60a28e64bd8 100644 --- a/policyengine_us/tests/core/_social_security_state_backend.py +++ b/policyengine_us/tests/core/_social_security_state_backend.py @@ -58,6 +58,7 @@ from policyengine_core.data_storage.store_history import ( StoreHistory, _period_from_storage_key, + advance_sequence_past, next_sequence_number, ) from policyengine_core.periods import Period @@ -741,6 +742,7 @@ def module(name, **values): json=json, ContextVar=ContextVar, WeakKeyDictionary=WeakKeyDictionary, + advance_sequence_past=advance_sequence_past, ref=ref, dataclass=dataclass, wraps=wraps, diff --git a/policyengine_us/tests/core/test_social_security_real_core.py b/policyengine_us/tests/core/test_social_security_real_core.py index 6c0d83361b8..ea049a0ca17 100644 --- a/policyengine_us/tests/core/test_social_security_real_core.py +++ b/policyengine_us/tests/core/test_social_security_real_core.py @@ -8,6 +8,7 @@ from __future__ import annotations import importlib +import itertools import pickle import sys import types @@ -647,15 +648,42 @@ def test_nonmemoized_store_keeps_the_existing_fast_return_date( @pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) -def test_pickled_fast_only_cache_retains_its_return_dates(mode): +def test_pickled_fast_only_cache_retains_its_return_dates(mode, monkeypatch): + from policyengine_core.data_storage import store_history + + stores = [] + original_record_store = Holder._record_store + + def record_store(holder, requested, number): + if holder.variable.name == "independent_formula": + stores.append(number) + return original_record_store(holder, requested, number) + + monkeypatch.setattr(Holder, "_record_store", record_store) simulation = fresh_without_holder_cache( mode, ["independent_formula"], {PUBLIC: {"2024": [12_000.0]}} ) assert scalar(simulation, PUBLIC + "_reported", "2024") == 12_000 assert scalar(simulation, "independent_formula", "2024") == 77 key = ("independent_formula", period("2024")) + history = simulation._get_store_history() + first_number = history.earliest_dependency(*key) + assert first_number == stores[-1] assert stored(simulation, *key) is None - simulation._fast_cache = pickle.loads(pickle.dumps(simulation._fast_cache)) + simulation._fast_cache.pop(key) + store_history.advance_sequence_past(store_history.next_sequence_number() + 1_000) + assert scalar(simulation, "independent_formula", "2024") == 77 + memo_number = stores[-1] + assert first_number < memo_number + assert history.earliest_dependency(*key) == first_number + blob = pickle.dumps(simulation._fast_cache) + # A different process can start below the restored memo's own store + # number. Restoring only StoreHistory's first number would be too early. + with monkeypatch.context() as counter_patch: + counter_patch.setattr(store_history, "_sequence", itertools.count(1)) + restored = pickle.loads(blob) + assert store_history.next_sequence_number() > memo_number + simulation._fast_cache = restored np.testing.assert_array_equal(simulation._fast_cache[key], [77]) assert stored(simulation, *key) is None epoch = simulation._input_epoch diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index b5abeeb83e4..af8a0a6f976 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -20,7 +20,10 @@ from pathlib import Path from weakref import WeakKeyDictionary, ref -from policyengine_core.data_storage.store_history import next_sequence_number +from policyengine_core.data_storage.store_history import ( + advance_sequence_past, + next_sequence_number, +) from policyengine_core.periods import period as to_period from policyengine_core.simulations.simulation import NpEncoder @@ -490,6 +493,8 @@ def forget(reference): def _restore_formula_fast_cache(values, numbers): + if numbers: + advance_sequence_past(max(numbers.values())) return _FormulaFastCache(values, numbers) From 5c0e93ecd68afb2e6fe95fe9d894087437c3bb8d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 18:54:44 -0700 Subject: [PATCH 61/76] Preserve formula dates across duplicate memo assignments --- .../tests/core/test_social_security_real_core.py | 2 ++ policyengine_us/tools/social_security_inputs.py | 6 +++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/policyengine_us/tests/core/test_social_security_real_core.py b/policyengine_us/tests/core/test_social_security_real_core.py index ea049a0ca17..8705e64760d 100644 --- a/policyengine_us/tests/core/test_social_security_real_core.py +++ b/policyengine_us/tests/core/test_social_security_real_core.py @@ -615,6 +615,8 @@ def test_nonmemoized_store_keeps_the_existing_fast_return_date( assert scalar(simulation, "independent_formula", "2023") == 77 key = ("independent_formula", period("2023")) memo = simulation._fast_cache[key] + # A calculation backend may repeat the insertion without another store. + simulation._fast_cache[key] = memo number = simulation._get_store_history().earliest_dependency(*key) if memo_before_since: simulation.set_input(PUBLIC, "2024", [12_000.0]) diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index af8a0a6f976..706085c25c8 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -513,8 +513,12 @@ def __init__(self, values=(), numbers=None): self._owner = None def __setitem__(self, key, array): + # Reassigning the same returned object without another formula store + # does not change its date (some calculation backends do this twice). + number = self._store_numbers.get(key) if self.get(key) is array else None super().__setitem__(key, array) - self._store_numbers.pop(key, None) + if number is None: + self._store_numbers.pop(key, None) simulation = None if self._owner is None else self._owner() if simulation is None: return From 6498687a0d59f57c736ada34a8d9861febe17128 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 21:56:10 -0700 Subject: [PATCH 62/76] Preserve formula return dates across bulk cache deletion --- changelog.d/ss-benefit-levers.added.md | 2 +- .../core/test_social_security_real_core.py | 195 ++++++++++++++++++ .../tools/social_security_inputs.py | 47 ++++- 3 files changed, 238 insertions(+), 6 deletions(-) diff --git a/changelog.d/ss-benefit-levers.added.md b/changelog.d/ss-benefit-levers.added.md index 2a7fe02d748..162ee4c65c9 100644 --- a/changelog.d/ss-benefit-levers.added.md +++ b/changelog.d/ss-benefit-levers.added.md @@ -1 +1 @@ -Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs remain independent sources, exposed through a read-time reported view without automatic twin arrays or provenance stamps. Initial loading routes sources so same-year benefit factors apply; later public writes remain final overrides. Explicit reported inputs retain priority, while exact compound totals keep their period semantics. Compound helpers retain initial-input precedence. Source refreshes use Core's dependency sequence cutoff, including formula returns kept only in the fast cache, preserving raw caller values, supplied inputs and earlier formula results. Builders, partial loading failures, copies, cache replay, extraction, dump restoration and subsampling preserve successful source inputs. A state machine searches combinations of write and copy routes. +Add composable Social Security benefit levers, pure reported component inputs, and a stored-data scoring preflight. Legacy public inputs remain independent sources, exposed through a read-time reported view without automatic twin arrays or provenance stamps. Initial loading routes sources so same-year benefit factors apply; later public writes remain final overrides. Explicit reported inputs retain priority, while exact compound totals keep their period semantics. Compound helpers retain initial-input precedence. Source refreshes use Core's dependency sequence cutoff, including formula returns kept only in the fast cache, preserving raw caller values, supplied inputs and earlier formula results. Bulk cache deletion preserves surviving formula dates; a formula return whose date is lost is conservatively recalculated after a dependent source write. Builders, partial loading failures, copies, cache replay, extraction, dump restoration and subsampling preserve successful source inputs. A state machine searches combinations of write and copy routes. diff --git a/policyengine_us/tests/core/test_social_security_real_core.py b/policyengine_us/tests/core/test_social_security_real_core.py index 8705e64760d..b29fb14af6f 100644 --- a/policyengine_us/tests/core/test_social_security_real_core.py +++ b/policyengine_us/tests/core/test_social_security_real_core.py @@ -530,6 +530,201 @@ def test_source_write_refreshes_fast_only_formula_returns(mode, route, source, i assert simulation._input_epoch > epoch +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +@pytest.mark.parametrize( + "route", ("dataset", "reported_write", "branch_reported_write") +) +def test_bulk_delete_keeps_fast_only_dates_for_later_source_refresh(mode, route): + simulation = fresh_without_holder_cache( + mode, [PUBLIC, PUBLIC + "_base"], {PUBLIC: {"2024": [12_000.0]}} + ) + if route == "branch_reported_write": + simulation = simulation.get_branch("child") + assert scalar(simulation, PUBLIC, "2024") == 12_000 + key = (PUBLIC, period("2024")) + before = simulation._fast_cache[key] + number = simulation._fast_cache._store_numbers[key] + assert number >= earliest_source_dependency(simulation) + assert stored(simulation, *key) is None + assert stored(simulation, PUBLIC + "_base", "2024") is None + simulation.delete_arrays("independent_formula") + assert simulation._fast_cache[key] is before + epoch = simulation._input_epoch + if route == "dataset": + rebuild_with_source(simulation, 24_000.0) + elif route == "reported_write": + simulation.set_input(PUBLIC + "_reported", "2024", [24_000.0]) + else: + simulation.get_holder(PUBLIC + "_reported")._set( + period("2024"), + np.array([24_000.0]), + simulation.branch_name, + is_input=True, + ) + memo_survived_write = key in simulation._fast_cache + reported_after = scalar(simulation, PUBLIC + "_reported", "2024") + public_after = scalar(simulation, PUBLIC, "2024") + assert (reported_after, public_after) == (24_000, 24_000) + assert not memo_survived_write + assert simulation._input_epoch > epoch + + +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +def test_bulk_delete_preserves_fast_only_return_before_since(mode): + simulation = fresh_without_holder_cache(mode, ["independent_formula"]) + assert scalar(simulation, "independent_formula", "2023") == 77 + key = ("independent_formula", period("2023")) + before = simulation._fast_cache[key] + number = simulation._fast_cache._store_numbers[key] + simulation.set_input(PUBLIC, "2024", [12_000.0]) + assert scalar(simulation, PUBLIC + "_reported", "2024") == 12_000 + assert number < earliest_source_dependency(simulation) + simulation.delete_arrays(PUBLIC + "_base") + assert simulation._fast_cache[key] is before + # The original return's own date must survive Core's dictionary replacement. + assert getattr(simulation._fast_cache, "_store_numbers", {}).get(key) == number + epoch = simulation._input_epoch + simulation.set_input(PUBLIC + "_reported", "2024", [24_000.0]) + assert simulation._fast_cache[key] is before + assert simulation._input_epoch == epoch + assert scalar(simulation, "independent_formula", "2023") == 77 + + +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +def test_bulk_delete_does_not_transfer_date_to_a_different_array(mode, monkeypatch): + simulation = fresh_without_holder_cache(mode, ["independent_formula"]) + assert scalar(simulation, "independent_formula", "2023") == 77 + key = ("independent_formula", period("2023")) + before = simulation._fast_cache[key] + replacement = before.copy() + assert key in simulation._fast_cache._store_numbers + original_delete = hooks._PHYSICAL["delete_arrays"] + + def replace_surviving_return(owner, name, requested=None): + result = original_delete(owner, name, requested) + # A backend can replace a return while rebuilding its dictionary. + # Equal values do not make this array the original numbered return. + owner._fast_cache[key] = replacement + return result + + monkeypatch.setitem(hooks._PHYSICAL, "delete_arrays", replace_surviving_return) + simulation.delete_arrays(PUBLIC + "_base") + assert simulation._fast_cache[key] is replacement + assert simulation._fast_cache[key] is not before + assert key not in simulation._fast_cache._store_numbers + + +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +def test_bulk_delete_component_removes_its_date_and_keeps_base_refreshable(mode): + simulation = fresh_without_holder_cache( + mode, [PUBLIC, PUBLIC + "_base"], {PUBLIC: {"2024": [12_000.0]}} + ) + assert scalar(simulation, PUBLIC, "2024") == 12_000 + key = (PUBLIC, period("2024")) + base_key = (PUBLIC + "_base", period("2024")) + base_before = simulation._fast_cache[base_key] + simulation.delete_arrays(PUBLIC) + assert key not in simulation._fast_cache + assert key not in getattr(simulation._fast_cache, "_store_numbers", {}) + assert simulation._fast_cache[base_key] is base_before + epoch = simulation._input_epoch + simulation.set_input(PUBLIC + "_reported", "2024", [24_000.0]) + assert scalar(simulation, PUBLIC, "2024") == 24_000 + assert simulation._input_epoch > epoch + + +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +@pytest.mark.parametrize("name", (PUBLIC, "independent_formula")) +def test_period_delete_keeps_surviving_fast_only_dates(mode, name): + simulation = fresh_without_holder_cache( + mode, + [PUBLIC, PUBLIC + "_base", "independent_formula"], + {PUBLIC: {"2024": [12_000.0]}}, + ) + assert scalar(simulation, PUBLIC, "2024") == 12_000 + assert scalar(simulation, "independent_formula", "2024") == 77 + key = (name, period("2024")) + survivor = ( + PUBLIC + "_base" if name == PUBLIC else PUBLIC, + period("2024"), + ) + before = simulation._fast_cache[survivor] + number = simulation._fast_cache._store_numbers[survivor] + simulation.delete_arrays(name, "2024") + assert key not in simulation._fast_cache + assert key not in simulation._fast_cache._store_numbers + assert simulation._fast_cache[survivor] is before + assert simulation._fast_cache._store_numbers[survivor] == number + epoch = simulation._input_epoch + simulation.set_input(PUBLIC + "_reported", "2024", [24_000.0]) + assert scalar(simulation, PUBLIC, "2024") == 24_000 + assert simulation._input_epoch > epoch + + +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +def test_undated_fast_only_return_is_intentionally_dropped_without_touching_callers( + mode, +): + names = [PUBLIC, PUBLIC + "_base", "independent_formula"] + simulation = fresh_without_holder_cache(mode, names, {PUBLIC: {"2024": [12_000.0]}}) + assert scalar(simulation, PUBLIC, "2024") == 12_000 + key = (PUBLIC, period("2024")) + holder = simulation.get_holder("independent_formula") + simulation.set_input("independent_formula", "2025", [99.0]) + holder._set( + period("2026"), + np.array([101.0]), + simulation.branch_name, + derived=False, + is_input=False, + ) + supplied_input = stored(simulation, "independent_formula", "2025") + raw_seed = stored(simulation, "independent_formula", "2026") + assert supplied_input is not None and raw_seed is not None + assert holder._is_input(period("2025"), simulation.branch_name) + assert not holder._is_input(period("2026"), simulation.branch_name) + assert not holder.is_derived(period("2025"), simulation.branch_name) + assert not holder.is_derived(period("2026"), simulation.branch_name) + assert earliest_source_dependency(simulation) is not None + simulation._fast_cache = dict(simulation._fast_cache) + epoch = simulation._input_epoch + simulation.get_holder(PUBLIC + "_reported")._set( + period("2024"), + np.array([24_000.0]), + simulation.branch_name, + is_input=True, + ) + # Intended fail-safe: with a cutoff, an undated formula return is treated + # as post-cutoff. Physical inputs and raw seeds remain outside that rule. + assert key not in simulation._fast_cache + assert simulation._input_epoch > epoch + assert scalar(simulation, PUBLIC, "2024") == 24_000 + assert stored(simulation, "independent_formula", "2025") is supplied_input + assert stored(simulation, "independent_formula", "2026") is raw_seed + assert holder._is_input(period("2025"), simulation.branch_name) + assert not holder._is_input(period("2026"), simulation.branch_name) + assert scalar(simulation, "independent_formula", "2025") == 99 + assert scalar(simulation, "independent_formula", "2026") == 101 + + +@pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) +def test_bulk_delete_unread_source_load_has_no_cutoff_and_keeps_fast_return(mode): + names = [PUBLIC, PUBLIC + "_base", PUBLIC + "_reported", "independent_formula"] + simulation = fresh_without_holder_cache(mode, names) + assert scalar(simulation, "independent_formula", "2023") == 77 + key = ("independent_formula", period("2023")) + before = simulation._fast_cache[key] + simulation.delete_arrays(PUBLIC + "_base") + assert earliest_source_dependency(simulation) is None + assert simulation._fast_cache[key] is before + epoch = simulation._input_epoch + rebuild_with_source(simulation, 24_000.0) + # No source has been read, so even an undated fast-only return is retained. + assert simulation._fast_cache[key] is before + assert simulation._input_epoch == epoch + assert scalar(simulation, PUBLIC, "2024") == 24_000 + + @pytest.mark.parametrize("mode", ("blacklist", "variables_to_drop")) def test_fast_only_formula_before_since_survives_without_a_drop_witness(mode): simulation = fresh_without_holder_cache(mode, ["independent_formula"]) diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 706085c25c8..3cbfc659f8d 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -566,7 +566,7 @@ def __reduce__(self): def _derived_cached_returns(simulation): - """Remember numbered formula returns without retaining an old dataset. + """Remember formula returns without retaining an old dataset. Core's dataset rebuild detaches holders but leaves their fast-cache returns. Weak references and sequence numbers let the bounded refresh @@ -622,8 +622,14 @@ def has_source(holder, period, branches): # produced this return. Skipped-storage formula returns need their # own number, captured by _cache_result/_record_store below. number = getattr(fast_cache, "_store_numbers", {}).get(key) - if number is not None: - entries[key] = (number, ref(array)) + # Intended fail-safe (refresh-scope addendum): Core's two memo + # insertion paths return supplied inputs only when physically + # stored. After the source-view exclusions above, a fast-only + # return is therefore calculated, including defaults/carry-over. + # An unknown date cannot establish that it predates the cutoff. + # Keep None here so loading snapshots also retain this witness; + # _purge_derived_arrays still does nothing without a cutoff. + entries[key] = (number, ref(array)) continue if holder._memory_storage.has(period, branch): storage = holder._memory_storage @@ -690,10 +696,10 @@ def _purge_derived_arrays(simulation, public, period): if owner is not simulation: continue for key, (number, reference) in tuple(entries.items()): - if number >= since: + if number is None or number >= since: array = reference() if array is not None and simulation._fast_cache.get(key) is array: - # A numbered formula return can outlive its holder, or + # A formula return can outlive its holder, or # Core can skip holder storage while memoizing the return. simulation._fast_cache.pop(key) dropped += 1 @@ -746,6 +752,7 @@ def install_social_security_input_compatibility(): recorded_array=CoreSimulation._get_recorded_input_array, calculate=CoreSimulation._calculate, cache_result=CoreSimulation._cache_result, + delete_arrays=CoreSimulation.delete_arrays, record_store=Holder._record_store, stored_input=helpers.get_stored_input, ) @@ -809,6 +816,36 @@ def cache_result(simulation, holder, array, period, input_state): CoreSimulation._cache_result = cache_result + @wraps(_PHYSICAL["delete_arrays"]) + def delete_arrays(simulation, variable, period=None): + if period is not None or not _enabled(simulation): + return _PHYSICAL["delete_arrays"](simulation, variable, period) + # Core's bulk delete replaces the memo dictionary while retaining + # other variables' arrays. Preserve each surviving return's actual + # date, rather than redating it from StoreHistory's first store. + # The snapshot adds only numbers and weak array references. + entries = { + key: (number, ref(simulation._fast_cache[key])) + for key, number in getattr( + simulation._fast_cache, "_store_numbers", {} + ).items() + if key in simulation._fast_cache + } + try: + return _PHYSICAL["delete_arrays"](simulation, variable, period) + finally: + cache = simulation._fast_cache + numbers = { + key: number + for key, (number, reference) in entries.items() + if reference() is not None and cache.get(key) is reference() + } + cache = _FormulaFastCache(cache, numbers) + cache._owner = ref(simulation) + simulation._fast_cache = cache + + CoreSimulation.delete_arrays = delete_arrays + @wraps(_PHYSICAL["stored_input"]) def get_stored_input(holder, period, branch_name): value = _PHYSICAL["stored_input"](holder, period, branch_name) From df0ec80c7e8de37c56b83be5f108468cf85a0152 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 22:04:33 -0700 Subject: [PATCH 63/76] Generate bulk deletion and source-write state transitions --- .../core/_social_security_state_backend.py | 1 + ...est_social_security_input_state_machine.py | 47 +++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/policyengine_us/tests/core/_social_security_state_backend.py b/policyengine_us/tests/core/_social_security_state_backend.py index 60a28e64bd8..369e18abc0f 100644 --- a/policyengine_us/tests/core/_social_security_state_backend.py +++ b/policyengine_us/tests/core/_social_security_state_backend.py @@ -608,6 +608,7 @@ def load(self): ) sim_methods = ( "set_input", + "delete_arrays", "_get_store_history", "_get_requested_variables", "_drop_values_that_may_depend_on", diff --git a/policyengine_us/tests/core/test_social_security_input_state_machine.py b/policyengine_us/tests/core/test_social_security_input_state_machine.py index eca7c39ce12..aabc2ebaabb 100644 --- a/policyengine_us/tests/core/test_social_security_input_state_machine.py +++ b/policyengine_us/tests/core/test_social_security_input_state_machine.py @@ -667,6 +667,53 @@ def cache_result( value, as_period(period), simulation.branch_name, derived=True ) + @rule(slot=TARGET, name=st.sampled_from(CACHE_NAMES)) + def bulk_delete_formula_cache(self, slot, name): + """Exercise Core's replacement of the memo dict around other returns.""" + simulation, _ = self.state(slot) + holder = simulation.get_holder(name) + visible = set(simulation._get_visible_branch_names()) + if any( + branch in visible and not storage.is_derived(period, branch) + for storage in (holder._memory_storage, holder._disk_storage) + if storage is not None + for branch, period in storage.get_known_branch_periods() + ): + # Delete formula caches, preserving this machine's caller source + # model. Core's delete API does not refresh dependent formulas + # after deleting caller inputs; source writes are separate rules. + return + simulation.delete_arrays(name) + + @rule( + component=COMPONENT, + name=st.sampled_from(CACHE_NAMES), + amount=st.integers(min_value=12_001, max_value=48_000), + ) + def bulk_delete_then_reported_write(self, component, name, amount): + """Compose a replaced memo with a later source write before other reads.""" + simulation = self.backend.Builder().build_from_variables( + self.backend.system, {component: {"2024": [12_000]}} + ) + model = Model() + model.write(component, False, "2024", 12_000, True, "builder", initial=True) + assert float(simulation.calculate(component, as_period("2024"))[0]) == 24_000 + simulation.delete_arrays(name) + # Initial public inputs are on reserved source branches, outside + # Core's visible-branch bulk deletion. Every generated name therefore + # deletes only formula caches here, keeping the caller model intact. + source = simulation.get_holder(f"{component}_reported").get_array( + as_period("2024"), simulation.branch_name + ) + assert source is not None and float(source[0]) == 12_000 + simulation.set_input(f"{component}_reported", "2024", [amount]) + model.write(component, True, "2024", amount, True, "after bulk delete") + # Add this state after the composition: warming every component + # beforehand would create unrelated physical formulas whose later + # deletion clears the memo, masking a lost fast-only return's date. + # Both invariants still check every state after every generated step. + self.append(simulation, model) + @rule( slot=TARGET, component=COMPONENT, From a6dcc92217ae60127edd588295684f332713ff1a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 26 Oct 2026 22:50:47 -0700 Subject: [PATCH 64/76] Retain shared benefit parameters in calculator reference --- .../tests/unit/test_ss_calculator_differential.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/policyengine_us/tests/unit/test_ss_calculator_differential.py b/policyengine_us/tests/unit/test_ss_calculator_differential.py index 2014996c4ba..4b4ad08e875 100644 --- a/policyengine_us/tests/unit/test_ss_calculator_differential.py +++ b/policyengine_us/tests/unit/test_ss_calculator_differential.py @@ -1,8 +1,10 @@ """The optional calculator leaves the existing model disconnected and unchanged. -The reference removes the new variables and parameter namespaces entirely. It -therefore detects an accidental dependency even when a new variable's default -value happens to be neutral. This is portable in source distributions and CI; +The reference removes calculator-only variables and parameter namespaces. Shared +benefit-lever parameters and statutory COLA history remain available to the +reported-benefit formulas. It detects an accidental calculator dependency even +when a new variable's default value happens to be neutral. This is portable in +source distributions and CI; it does not need a second checkout or git at test time. The implementation's separate source-diff check establishes that legacy formulas/parameters were not edited. These comparisons exercise their live SSA and downstream dependency @@ -21,12 +23,12 @@ CALCULATOR_PREFIXES = ("ss_calculator_", "ss_calculated_") NEW_PARAMETER_NODES = ( - "gov.contrib.ssa.benefit_levers", + # Benefit levers and COLA history also serve the reported-benefit route, + # which remains part of the reference when the calculator is removed. "gov.contrib.ssa.pia", "gov.contrib.congress.s770", "gov.ssa.social_security.pia.bend_points", "gov.ssa.social_security.cola", - "gov.ssa.cola_history", ) YEAR = 2026 From a960ce6935c2ca3b61e708d68467e300a03c07d8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 27 Oct 2026 07:29:20 -0700 Subject: [PATCH 65/76] Add configurable prior-year Social Security COLA means test --- changelog.d/ss-cola-means-test-e3.added.md | 1 + docs/social-security-benefit-levers.md | 6 +- docs/social-security-cola-means-test.md | 243 +++++ docs/usage/social-security-calculator.md | 26 +- .../below_threshold_cola_index.yaml | 9 + .../cola/means_test/effective_month.yaml | 10 + .../cola/means_test/effective_year.yaml | 10 + .../cola/means_test/enabled.yaml | 10 + .../filing_groups/HEAD_OF_HOUSEHOLD.yaml | 9 + .../cola/means_test/filing_groups/JOINT.yaml | 9 + .../means_test/filing_groups/SEPARATE.yaml | 9 + .../cola/means_test/filing_groups/SINGLE.yaml | 9 + .../filing_groups/SURVIVING_SPOUSE.yaml | 9 + .../cola/means_test/first_eligible_year.yaml | 7 + .../cola/means_test/income_sensitivity.yaml | 9 + .../means_test/pre_receipt_cola_index.yaml | 9 + .../benefit_levers/cola/means_test/scope.yaml | 9 + .../cola/means_test/thresholds/base_year.yaml | 10 + .../cola/means_test/thresholds/index_lag.yaml | 10 + .../cola/means_test/thresholds/joint.yaml | 10 + .../cola/means_test/thresholds/rounding.yaml | 10 + .../cola/means_test/thresholds/single.yaml | 10 + .../reforms/ssa/cola_means_test.py | 40 + .../social_security_cola_means_test.yaml | 44 + .../tests/unit/test_ss_cola_means_test.py | 852 ++++++++++++++++++ .../tools/social_security/engine.py | 82 +- .../tools/social_security/means_test.py | 399 ++++++++ .../tools/social_security_benefit_levers.py | 78 +- .../social_security_cola_means_test_magi.py | 45 + ...ocial_security_cola_means_test_withheld.py | 14 + .../gov/ssa/ss/social_security_dependents.py | 20 +- ...ocial_security_dependents_cola_withheld.py | 6 +- .../gov/ssa/ss/social_security_disability.py | 20 +- ...ocial_security_disability_cola_withheld.py | 6 +- .../gov/ssa/ss/social_security_retirement.py | 20 +- ...ocial_security_retirement_cola_withheld.py | 6 +- .../gov/ssa/ss/social_security_survivors.py | 20 +- ...social_security_survivors_cola_withheld.py | 6 +- 38 files changed, 2027 insertions(+), 75 deletions(-) create mode 100644 changelog.d/ss-cola-means-test-e3.added.md create mode 100644 docs/social-security-cola-means-test.md create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/below_threshold_cola_index.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/effective_month.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/effective_year.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/enabled.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/HEAD_OF_HOUSEHOLD.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/JOINT.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/SEPARATE.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/SINGLE.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/SURVIVING_SPOUSE.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/first_eligible_year.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/income_sensitivity.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/pre_receipt_cola_index.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/scope.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/base_year.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/index_lag.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/joint.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/rounding.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/single.yaml create mode 100644 policyengine_us/reforms/ssa/cola_means_test.py create mode 100644 policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_cola_means_test.yaml create mode 100644 policyengine_us/tests/unit/test_ss_cola_means_test.py create mode 100644 policyengine_us/tools/social_security/means_test.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_cola_means_test_magi.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_cola_means_test_withheld.py diff --git a/changelog.d/ss-cola-means-test-e3.added.md b/changelog.d/ss-cola-means-test-e3.added.md new file mode 100644 index 00000000000..9b79cb1d69d --- /dev/null +++ b/changelog.d/ss-cola-means-test-e3.added.md @@ -0,0 +1 @@ +Add an opt-in Social Security COLA means test using prior-year MAGI excluding Social Security, configurable indexed thresholds and filing groups, cohort or all-beneficiary scope, statutory or chained CPI-U COLAs, income-history sensitivities, and an A9 preset. Apply each withholding decision jointly with other COLA changes in reported benefits and the household calculator, with current-law defaults unchanged. diff --git a/docs/social-security-benefit-levers.md b/docs/social-security-benefit-levers.md index 935fcfe4611..c09263434b6 100644 --- a/docs/social-security-benefit-levers.md +++ b/docs/social-security-benefit-levers.md @@ -115,8 +115,10 @@ For each exposed payment event, statutory baseline rate c and configured addition delta give adjusted rate max(0,c+delta), without a second rounding of the policy addition. The joint multiplier is `(1+adjusted)/(1+c)` unless that event is withheld, in which case it is `1/(1+c)`. Each component's -`_cola_withheld` term defaults false; a subsequent means-test implementation -supplies the rules. Switching a policy off later preserves its earlier effect +`_cola_withheld` term defaults false; the optional +[COLA means test](social-security-cola-means-test.md) supplies prior-year income +decisions when enabled. Explicit decisions remain available. Switching a policy +off later preserves its earlier effect on the benefit level. Public reported-route benefits multiply this joint factor, the PIA ratio factor, the across-the-board factor, and the payable factor once. diff --git a/docs/social-security-cola-means-test.md b/docs/social-security-cola-means-test.md new file mode 100644 index 00000000000..5c6e19e62cd --- /dev/null +++ b/docs/social-security-cola-means-test.md @@ -0,0 +1,243 @@ +# Social Security COLA means test + +The contributed COLA means test is disabled under current law. It can withhold +an annual COLA based on the beneficiary's tax unit, while retaining the benefit +level reached before that event. It supports both reported benefit components +and the supplied-record household calculator. Scenario thresholds and dates +must be supplied before scoring; PE-US does not choose package values. + +## Policy basis and choices + +The income test follows [OACT provision A9 on the 2026 Trustees Report +basis](https://www.ssa.gov/OACT/solvency/provisions/cola.html). Its instruction +is: “Use prior tax year income data for this determination.” The live page +was available when this implementation was prepared. + +A9 uses a chained CPI-U COLA below the income line and no COLA above it, for +all beneficiaries beginning with the December 2028 COLA. Its single, +head-of-household and separate filers share a line; joint filers have twice +that line. It also specifies chained CPI-U before benefit receipt. Its example +lines are $85,000 and $170,000 in 2018 dollars, indexed by CPI-U. The preset +below implements these mechanics. OACT's displayed December 2028 dollar +estimates use its projections; this model calculates indexed lines from its +own pinned CPI-U path. + +The default enabled scenario differs from A9 in two parameterized ways: it +retains the statutory COLA below the line and restricts the policy to people +first eligible from 2026. The effective date, first-eligible year, filing-group +mapping, dollar lines, base year, index lag and rounding increment are also +parameters. A surviving spouse maps to the joint line by default; this is a +modeling choice because A9 does not name that status. Only income strictly +above the line withholds a COLA. Equality receives the COLA, an explicit +convention for A9's otherwise unspecified equality case. + +## Parameters and activation + +All paths in this table are relative to +`gov.contrib.ssa.benefit_levers.cola.means_test`. +Choice parameters use numeric codes because the parameter engine does not +accept string-valued choices. The semantic names below label those codes; +reforms must supply the numbers. + +| Parameter | Current-law/default value | Scenario meaning | +| --- | --- | --- | +| `enabled` | `false` | Opt into the means test. | +| `scope` | `0` = `first_eligible_cohorts` | Alternatively, `1` = `all_beneficiaries`. | +| `first_eligible_year` | `2026` | Inclusive first-eligibility boundary in cohort mode. | +| `effective_year`, `effective_month` | `2026`, `12` | First affected effective COLA; ordinarily first paid January 2027. | +| `below_threshold_cola_index` | `0` = `statutory` | Alternatively, `1` = `chained_cpi_u`. | +| `pre_receipt_cola_index` | `-1` = `follow_below_threshold` | Follow the below-line choice, or explicitly select `0` = `statutory` or `1` = `chained_cpi_u`. | +| `income_sensitivity` | `0` = `prior_year` | Alternatively, `1` = `never_withheld` or `2` = `always_withheld`. | +| `thresholds.base_year` | `2026` | Dollar year of the supplied lines. | +| `thresholds.single`, `thresholds.joint` | `0`, `0` | Neutral placeholders; activation requires positive scenario lines. | +| `thresholds.index_lag` | `1` | Years before the payment year for the December CPI-U observation. | +| `thresholds.rounding` | `100` | Dollar increment, rounded half up. | +| `filing_groups.SINGLE` | `0` = `single` | Select the single line; `1` selects joint. | +| `filing_groups.HEAD_OF_HOUSEHOLD` | `0` = `single` | Select the single line; `1` selects joint. | +| `filing_groups.SEPARATE` | `0` = `single` | Select the single line; `1` selects joint. | +| `filing_groups.JOINT` | `1` = `joint` | Select the joint line; `0` selects single. | +| `filing_groups.SURVIVING_SPOUSE` | `1` = `joint` | Modeling choice; `0` selects single. | + +`means_test.enabled` activates both the income decision and the selected index +without requiring another switch. The enclosing `cola.withholding_enabled` +remains the A5 switch for externally resolved withholding decisions when the +means test is disabled. Explicit component/year values can override the +calculated decisions. + +The `a9` preset in `policyengine_us.reforms.ssa.cola_means_test` enables the +policy and selects: + +```yaml +scope: 1 # all_beneficiaries +effective_year: 2028 +effective_month: 12 +below_threshold_cola_index: 1 # chained_cpi_u +pre_receipt_cola_index: 1 # chained_cpi_u +income_sensitivity: 0 # prior_year +thresholds: + base_year: 2018 + single: 85000 + joint: 170000 + index_lag: 1 + rounding: 100 +``` + +Apply it as a reform, for example: + +```python +from policyengine_us import Simulation +from policyengine_us.reforms.ssa.cola_means_test import a9 + +simulation = Simulation(situation=situation, reform=a9) +``` + +These are public A9 example settings, not a scenario's package values. The +first affected checks are in January 2029. The ordinary contributed scenario +instead supplies its own positive 2026-dollar lines and retains the default +statutory/cohort choices unless its policy specifies otherwise. + +## Prior-year income and filing status + +`social_security_cola_means_test_magi` is a dedicated annual TaxUnit variable. +Its definition is AGI less taxable Social Security, plus tax-exempt interest +and the IRMAA-style exclusions under Internal Revenue Code sections 135, 911, +931 and 933. AGI retains its applicable deductions. This implements the +required MAGI excluding Social Security, replacing the earlier gross-source +income proposal. The [A9-linked actuarial memorandum, provision 5 on page +8](https://www.ssa.gov/OACT/solvency/SJohnson_20161208.pdf) connects the lines +to Medicare IRMAA; [H.R. 6489 section +103(d)](https://www.congress.gov/bill/114th-congress/house-bill/6489/text) +supplies the corresponding IRMAA definition reference. Excluding Social +Security and using the prior year are requirements of this contributed lever. + +For COLAs first paid in year `y`, the decision reads this MAGI and filing +status for tax year `y - 1`. It never substitutes current-year AGI. A supplied +MAGI history can directly set the dedicated variable for the relevant tax +years; supplied histories of its income sources and filing status also take +precedence over projected records. A tax unit's decision maps to each +beneficiary, subject to that component's exposure and cohort scope. + +Although prior-year AGI can depend on prior-year Social Security, every +benefit-related lookback moves to an earlier year. The dependency chain stops +before the policy's first affected event. It cannot return to the same +benefit/year pair. Multi-year dependency-graph tests exercise this decreasing +year chain, including deductions whose MAGI depends on taxable benefits. +The household engine resolves decisions only through the requested check +year for payment outputs. It does not eagerly read the income for the next +January COLA while computing this year's checks. December entitlement and +annual entitlement can require that next event; their prior-income dependency +uses the separate public payment route. + +Population scoring uses a **static-cross-section prior-income proxy**. +Dataset extension starts each year from the same original record and uprates +its monetary sources separately. The `y - 1` record is that same record's +extended prior-year copy, not an observed longitudinal income history. Ages, +cohorts and household composition are not carried forward as a person's +history. Supplied histories override this proxy. Compare +`income_sensitivity = 1` (`never_withheld`) and `2` (`always_withheld`) to +report the missing-income-history sensitivity envelope; both retain policy +timing and cohort scope. + +## Threshold indexing and COLA construction + +For filing group `g`, base year `b`, payment year `y`, index lag `L` and +rounding increment `R`, the income line is: + +```text +line(g, y) = round_half_up_to_R( + base_line(g) * CPI_U(December, y - L) / CPI_U(December, b) +) +withheld(y) = prior_year_MAGI > line(g, y) +``` + +The CPI-U series is `gov.bls.cpi.cpi_u`. The default observes the December +before the check year, including December 2028 for January 2029 checks. +Parameter snapshots are held until the next authored observation, including +in projections. The December snapshot, one-year payment lag and nearest-$100 +half-up rounding are documented modeling conventions: A9 specifies CPI-U +indexing without prescribing these details. Pin both the index path and +these settings for a score. + +The chained option uses Q3 averages of July, August and September levels from +`gov.bls.cpi.c_cpi_u`. As with the statutory CPI-W schedule, positive growth +is rounded half up to 0.1 percentage point; nonpositive or rounded-zero growth +gives no COLA. The comparison base is the Q3 associated with the last positive +index-derived COLA. A withheld year still advances that index schedule's +comparison base when its rounded rate is positive: income withholding must +not create a catch-up COLA later. Neither withholding nor a downward rate +adjustment reduces the nominal benefit level. Historical benefit-rounding +rules remain attached to their individual events. + +Before receipt, `pre_receipt_cola_index` controls COLAs accrued since initial +eligibility. Its default follows `below_threshold_cola_index`; the A9 preset +therefore uses chained CPI-U for these events. From the policy's effective +start, these pre-receipt events use the chosen index without an income test. +Supply entitlement and worker eligibility dates when known. Missing-date +cohort and exposure approximations remain those described in the +[benefit-lever guide](social-security-benefit-levers.md), including +adjacent-cohort sensitivities for integer-age uncertainty. + +An active A5 `cola.alternative_rate` overrides the index selected here. The +applicable `cola.pp_change` is then added to that rate, and the final rate is +floored at zero. Withholding overrides the entire resulting increase. + +## One joint COLA adjustment + +Each annual Person variable `social_security_X_cola_withheld`, for `X` equal +to `retirement`, `disability`, `survivors` or `dependents`, has a formula that +resolves its tax unit's decision. Explicit values override that formula. +For a reported benefit, the population adjustment accumulates one ratio for +each exposed event: + +```text +event numerator = 1 if withheld + 1 + adjusted_COLA otherwise +event ratio = event numerator / (1 + statutory_COLA) +joint factor = product of the exposed event ratios +``` + +The adjusted rate includes the selected index and applicable percentage-point +change, with a floor of zero. For example, statutory rates of 2% and 3% plus a +CPI-E-style addition of 0.15 percentage points give numerators `1.0215` and +`1.0315` if paid. If the second year is withheld, the two-year factor is +`(1.0215 / 1.02) * (1 / 1.03)`, approximately `0.97230`. The reformed level +from a pre-event amount of $1,000 is $1,021.50, compared with the statutory +$1,050.60. The withheld event receives no CPI-E increase. An independently +multiplied withholding factor would give the wrong result. + +The household calculator applies the same per-year decisions inside its COLA +schedule. It rounds the resulting PIA under the applicable event rule and +applies claiming and payment calculations once. Its explicit annual +`ss_calculator_withhold_cola` value takes precedence over the shared decision. + +For a public retirement component, select `social_security_retirement_source` +as `supplied_aime` or `supplied_pia`, and supply the matching +`ss_calculator_source` of `aime` or `pia`. For a restricted survivor record, +select `social_security_survivors_source = supplied_ratio_survivor` and +`ss_calculator_source = survivor`. These routes retain the calculator's input +and scope requirements in the [calculator +guide](usage/social-security-calculator.md). + +Calculated public routes bypass all population benefit factors: the engine +already applies the COLA, across-the-board and payable policies to its ledger. +Reported sources use the population factors. This avoids applying any of +these policies twice. Switching a policy off later does not restore an +earlier withheld increase. + +## Validation contract + +Validation covers all statutory/chained and cohort/all-beneficiary +combinations, both sides of the first-eligibility boundary, the A9 preset, +all filing statuses, equality, supplied histories, and both income +sensitivities. Boundary cases include a zero statutory COLA, negative rate +adjustments, high-water years and combined CPI-E changes and withholding. +Multi-year graph checks establish that the prior-year income read has no +cycle. Calculated-source cases check that the COLA policy is applied once. + +Property checks require disabled-policy equality to the stacked base, +nonincreasing benefits when a year is withheld, monotone withholding as income +rises at a fixed line, and equality of the joint factor to the product of its +per-year ratios. Differential checks cover reported situations, downstream +outputs and household calculator examples with this means test disabled. +Population impact requires a separately pinned scoring run; these mechanics +do not substitute OACT's solvency estimates for a PE-US result. diff --git a/docs/usage/social-security-calculator.md b/docs/usage/social-security-calculator.md index adca67d9941..8f9ce8fd72a 100644 --- a/docs/usage/social-security-calculator.md +++ b/docs/usage/social-security-calculator.md @@ -220,23 +220,23 @@ Withholding sets that event's COLA to zero, jointly with any rate change; no second COLA factor is applied. Payable shares must lie in [0,1], and rates must not imply a negative benefit. A depletion-year midyear blend is supplied in the annual payable-share path. E2 and E3 income/threshold computation belong to -their respective modules; their policies and jointly calculated tax-unit -outputs can consume this route once public wiring exists. +their respective modules. The [COLA means test](../social-security-cola-means-test.md) +now supplies these decisions from prior-tax-year MAGI and filing status. ## Public source-selector wiring -Source-selector PR #10400 was open when this route was built. Its canonical -COLA history and common lever parameter files are reused here. Historical bend +The source selectors reuse the canonical COLA history and common lever +parameter files. Historical bend points remain in a separate table because copying its legacy scale backfill -would change existing pre-2015 `ss_pia` outputs with this route disabled. This change adds -no `_reported` twins and changes no public benefit formula. Once it merges, -map its supplied-AIME/PIA retirement modes to `ss_calculated_retirement` and its -restricted survivor mode to `ss_calculated_survivors`. Bypass population -PIA/COLA factors and any across/payable factors already applied in this ledger; -perform the earnings test once. Preserve explicit public final overrides. -For joint household tax calculations before that wiring, explicitly supply the -calculated annual checks to the existing public components in a separate -calculation; never write them into `_reported` as if they were original data. +would change existing pre-2015 `ss_pia` outputs with this route disabled. +Select `social_security_retirement_source` as `supplied_aime` or `supplied_pia` +with matching `ss_calculator_source` of `aime` or `pia`. For the restricted +survivor route, select `social_security_survivors_source` as +`supplied_ratio_survivor` with calculator source `survivor`. The public benefit +then reads the calculated checks, bypassing population PIA, COLA, across-the-board +and payable factors already applied in this ledger. The earnings test runs +once. Explicit public final overrides retain their precedence. Calculated +checks are never written into `_reported` as original data. ## Sources and validation diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/below_threshold_cola_index.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/below_threshold_cola_index.yaml new file mode 100644 index 00000000000..bf563f666c3 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/below_threshold_cola_index.yaml @@ -0,0 +1,9 @@ +description: "COLA index for income at or below the means-test threshold: 0 = statutory pays the full statutory COLA; 1 = chained_cpi_u uses gov.bls.cpi.c_cpi_u with statutory rounding, high-water and no-decrease rules. The A9 preset uses 1. Numeric codes provide compatibility with the parameter engine." +values: + 0000-01-01: 0 +metadata: + label: Social Security below-threshold COLA index + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/effective_month.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/effective_month.yaml new file mode 100644 index 00000000000..7a1279b3b56 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/effective_month.yaml @@ -0,0 +1,10 @@ +description: First effective COLA calendar month for the contributed means test. The default December COLA first changes January payments in the following calendar year. +values: + 0000-01-01: 12 +metadata: + unit: month + label: Social Security COLA means-test effective month + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/effective_year.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/effective_year.yaml new file mode 100644 index 00000000000..0844cc81316 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/effective_year.yaml @@ -0,0 +1,10 @@ +description: First effective COLA calendar year for the contributed means test. A December COLA first changes January payments in the following year. The A9 preset uses December 2028. +values: + 0000-01-01: 2026 +metadata: + unit: year + label: Social Security COLA means-test effective year + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/enabled.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/enabled.yaml new file mode 100644 index 00000000000..f511d99f6cf --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/enabled.yaml @@ -0,0 +1,10 @@ +description: Enable the contributed Social Security COLA means test. Activation requires positive single and joint income thresholds supplied by the scenario. +values: + 0000-01-01: false +metadata: + unit: bool + label: Social Security COLA means test enabled + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/HEAD_OF_HOUSEHOLD.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/HEAD_OF_HOUSEHOLD.yaml new file mode 100644 index 00000000000..67cc10e808f --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/HEAD_OF_HOUSEHOLD.yaml @@ -0,0 +1,9 @@ +description: "Income threshold group for a head of household under the COLA means test: 0 = single; 1 = joint. OACT A9 uses 0; scenarios may supply another mapping." +values: + 0000-01-01: 0 +metadata: + label: Social Security COLA means-test head of household group + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/JOINT.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/JOINT.yaml new file mode 100644 index 00000000000..04c1be6c5ed --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/JOINT.yaml @@ -0,0 +1,9 @@ +description: "Income threshold group for joint filers under the COLA means test: 0 = single; 1 = joint. OACT A9 uses 1; scenarios may supply another mapping." +values: + 0000-01-01: 1 +metadata: + label: Social Security COLA means-test joint filer group + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/SEPARATE.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/SEPARATE.yaml new file mode 100644 index 00000000000..3e0e554b285 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/SEPARATE.yaml @@ -0,0 +1,9 @@ +description: "Income threshold group for married filing separately under the COLA means test: 0 = single; 1 = joint. OACT A9 uses 0; scenarios may supply another mapping." +values: + 0000-01-01: 0 +metadata: + label: Social Security COLA means-test separate filer group + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/SINGLE.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/SINGLE.yaml new file mode 100644 index 00000000000..d387884fe99 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/SINGLE.yaml @@ -0,0 +1,9 @@ +description: "Income threshold group for a single filer under the COLA means test: 0 = single; 1 = joint. OACT A9 uses 0; scenarios may supply another mapping." +values: + 0000-01-01: 0 +metadata: + label: Social Security COLA means-test single filer group + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/SURVIVING_SPOUSE.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/SURVIVING_SPOUSE.yaml new file mode 100644 index 00000000000..b556c9f6558 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/filing_groups/SURVIVING_SPOUSE.yaml @@ -0,0 +1,9 @@ +description: "Income threshold group for a surviving spouse under the COLA means test: 0 = single; 1 = joint. OACT A9 does not name this filing status; the default 1 is an explicit modeling choice that scenarios may change." +values: + 0000-01-01: 1 +metadata: + label: Social Security COLA means-test surviving spouse group + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/first_eligible_year.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/first_eligible_year.yaml new file mode 100644 index 00000000000..1288983bef2 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/first_eligible_year.yaml @@ -0,0 +1,7 @@ +description: First eligibility cohort reached when the Social Security COLA means-test scope is first_eligible_cohorts. This contributed reform is disabled by default. +values: + 0000-01-01: 2026 +metadata: + unit: year + label: Social Security COLA means-test first eligibility cohort + economy: false diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/income_sensitivity.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/income_sensitivity.yaml new file mode 100644 index 00000000000..094768cab92 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/income_sensitivity.yaml @@ -0,0 +1,9 @@ +description: "Income decision for the COLA means test: 0 = prior_year uses the preceding tax year's MAGI excluding Social Security; 1 = never_withheld and 2 = always_withheld provide the missing-history sensitivity envelope. Extended population data are static cross-sections, so the prior-year copy of a record is a proxy rather than an observed income history; supplied histories override that proxy. Numeric codes provide compatibility with the parameter engine." +values: + 0000-01-01: 0 +metadata: + label: Social Security COLA means-test income sensitivity + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/pre_receipt_cola_index.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/pre_receipt_cola_index.yaml new file mode 100644 index 00000000000..3c70485234e --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/pre_receipt_cola_index.yaml @@ -0,0 +1,9 @@ +description: "COLA index before benefit receipt: -1 = follow_below_threshold uses below_threshold_cola_index; 0 = statutory and 1 = chained_cpi_u override it. The default therefore follows the below-threshold policy. The A9 preset uses 1. Numeric codes provide compatibility with the parameter engine." +values: + 0000-01-01: -1 +metadata: + label: Social Security pre-receipt COLA index + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/scope.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/scope.yaml new file mode 100644 index 00000000000..ef35ac91277 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/scope.yaml @@ -0,0 +1,9 @@ +description: "COLA means-test scope: 0 = first_eligible_cohorts limits the reform to the configured first eligibility cohort and later; 1 = all_beneficiaries applies to every beneficiary. The A9 preset uses 1. Numeric codes provide compatibility with the parameter engine." +values: + 0000-01-01: 0 +metadata: + label: Social Security COLA means-test scope + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/base_year.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/base_year.yaml new file mode 100644 index 00000000000..2a4dffb9efc --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/base_year.yaml @@ -0,0 +1,10 @@ +description: Dollar base year for supplied single and joint COLA means-test thresholds. Scenario values are supplied at scoring time; the A9 example is 85000 single and 170000 joint in 2018 dollars. +values: + 0000-01-01: 2026 +metadata: + unit: year + label: Social Security COLA means-test threshold dollar base year + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/index_lag.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/index_lag.yaml new file mode 100644 index 00000000000..7987aa45ba8 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/index_lag.yaml @@ -0,0 +1,10 @@ +description: Years subtracted from a COLA payment year to select the December CPI-U threshold-index snapshot, relative to December of the dollar base year. The default uses the preceding December. This lag is a modeling convention; OACT A9 specifies CPI-U indexing without this snapshot rule. +values: + 0000-01-01: 1 +metadata: + unit: year + label: Social Security COLA means-test threshold CPI-U lag + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/joint.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/joint.yaml new file mode 100644 index 00000000000..af7e663bc01 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/joint.yaml @@ -0,0 +1,10 @@ +description: Joint-group COLA means-test income threshold in base-year dollars. Zero is an unset sentinel; activation requires a positive scenario value. The documented A9 example is 170000 in 2018 dollars, not a default package value. +values: + 0000-01-01: 0 +metadata: + unit: currency-USD + label: Social Security COLA means-test joint income threshold + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/rounding.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/rounding.yaml new file mode 100644 index 00000000000..16ceea795e4 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/rounding.yaml @@ -0,0 +1,10 @@ +description: Dollar increment for half-up rounding of CPI-U-indexed COLA means-test thresholds. The example increment of 100 dollars is a modeling convention reflecting the precision of OACT A9's displayed estimates; it is configurable for scenarios. +values: + 0000-01-01: 100 +metadata: + unit: currency-USD + label: Social Security COLA means-test threshold rounding increment + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/single.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/single.yaml new file mode 100644 index 00000000000..73bd9c7d326 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/single.yaml @@ -0,0 +1,10 @@ +description: Single-group COLA means-test income threshold in base-year dollars. Zero is an unset sentinel; activation requires a positive scenario value. The documented A9 example is 85000 in 2018 dollars, not a default package value. +values: + 0000-01-01: 0 +metadata: + unit: currency-USD + label: Social Security COLA means-test single income threshold + economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/reforms/ssa/cola_means_test.py b/policyengine_us/reforms/ssa/cola_means_test.py new file mode 100644 index 00000000000..8252a4b012a --- /dev/null +++ b/policyengine_us/reforms/ssa/cola_means_test.py @@ -0,0 +1,40 @@ +"""Named OACT A9 mechanics; applying this preset opts into the contribution.""" + +from policyengine_core.reforms import Reform + + +class a9(Reform): + def apply(self): + self.modify_parameters(self.configure_parameters) + + @staticmethod + def configure_parameters(parameters): + p = parameters.gov.contrib.ssa.benefit_levers.cola.means_test + values = { + "enabled": True, + "scope": 1, + "effective_year": 2028, + "effective_month": 12, + "below_threshold_cola_index": 1, + "pre_receipt_cola_index": 1, + "income_sensitivity": 0, + } + for name, value in values.items(): + getattr(p, name).update(start="0001-01-01", value=value) + for name, value in { + "base_year": 2018, + "single": 85000, + "joint": 170000, + "index_lag": 1, + "rounding": 100, + }.items(): + getattr(p.thresholds, name).update(start="0001-01-01", value=value) + for name, value in { + "SINGLE": 0, + "JOINT": 1, + "SEPARATE": 0, + "HEAD_OF_HOUSEHOLD": 0, + "SURVIVING_SPOUSE": 1, + }.items(): + getattr(p.filing_groups, name).update(start="0001-01-01", value=value) + return parameters diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_cola_means_test.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_cola_means_test.yaml new file mode 100644 index 00000000000..3cf0c634490 --- /dev/null +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_cola_means_test.yaml @@ -0,0 +1,44 @@ +- name: Prior-year income withholds only above the line in covered cohorts + period: 2027 + absolute_error_margin: 0.01 + input: + gov.contrib.ssa.benefit_levers.cola.means_test.enabled: true + gov.contrib.ssa.benefit_levers.cola.means_test.thresholds.single: 100000 + gov.contrib.ssa.benefit_levers.cola.means_test.thresholds.joint: 200000 + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + people: + equal: + age: 63 + social_security_retirement_worker_eligibility_year: 2026 + social_security_retirement_first_cola_year: 2027 + social_security_retirement_reported: 12000 + above: + age: 63 + social_security_retirement_worker_eligibility_year: 2026 + social_security_retirement_first_cola_year: 2027 + social_security_retirement_reported: 12000 + older_cohort: + age: 64 + social_security_retirement_worker_eligibility_year: 2025 + social_security_retirement_first_cola_year: 2027 + social_security_retirement_reported: 12000 + tax_units: + equal: + members: [equal] + social_security_cola_means_test_magi: + 2026: 100000 + above: + members: [above] + social_security_cola_means_test_magi: + 2026: 100001 + older_cohort: + members: [older_cohort] + social_security_cola_means_test_magi: + 2026: 100001 + output: + social_security_retirement_cola_withheld: [false, true, false] + # January 2027 statutory rate is 3.6%. The CPI-E addition appears only + # in the two paid numerators; the withheld numerator is exactly one. + social_security_retirement: + [12017.374517374517, 11583.011583011583, 12017.374517374517] diff --git a/policyengine_us/tests/unit/test_ss_cola_means_test.py b/policyengine_us/tests/unit/test_ss_cola_means_test.py new file mode 100644 index 00000000000..6c24e640e6a --- /dev/null +++ b/policyengine_us/tests/unit/test_ss_cola_means_test.py @@ -0,0 +1,852 @@ +"""COLA means-test acceptance, historical inputs and joint-event invariants. + +Synthetic thresholds are examples, never contributed package settings. The +payment-year convention means December2026's COLA first changes January2027. +Country systems are reused; pure schedule properties do not clone the model. +""" + +from copy import deepcopy +from decimal import Decimal, ROUND_DOWN +from math import prod + +import numpy as np +import pytest +from hypothesis import given, settings, strategies as st +from policyengine_core.parameters import ParameterNode + +from policyengine_us import Simulation +from policyengine_us.reforms.ssa.cola_means_test import a9 +from policyengine_us.tools.social_security.means_test import ( + chained_cola_rates, + indexed_threshold, +) +from policyengine_us.tools.social_security_benefit_levers import ( + cola_event_multiplier, + cola_events, + joint_cola_factor, +) + + +MEANS_TEST = "gov.contrib.ssa.benefit_levers.cola.means_test" +COMPONENTS = ("retirement", "disability", "survivors", "dependents") + + +def _reform(**overrides): + values = { + "enabled": True, + "scope": "first_eligible_cohorts", + "first_eligible_year": 2026, + "effective_year": 2026, + "effective_month": 12, + "below_threshold_cola_index": "statutory", + "thresholds.base_year": 2026, + "thresholds.single": 100_000, + "thresholds.joint": 200_000, + "thresholds.index_lag": 1, + "thresholds.rounding": 100, + } + values.update(overrides) + enum_values = { + "scope": {"first_eligible_cohorts": 0, "all_beneficiaries": 1}, + "below_threshold_cola_index": {"statutory": 0, "chained_cpi_u": 1}, + "pre_receipt_cola_index": { + "follow_below_threshold": -1, + "statutory": 0, + "chained_cpi_u": 1, + }, + "income_sensitivity": { + "prior_year": 0, + "never_withheld": 1, + "always_withheld": 2, + }, + } + for name, value in tuple(values.items()): + mapping = ( + {"single": 0, "joint": 1} + if name.startswith("filing_groups.") + else enum_values.get(name, {}) + ) + values[name] = mapping.get(value, value) + return { + f"{MEANS_TEST}.{name}": {"2020-01-01": value} for name, value in values.items() + } + + +def _record(eligibility=2026, **overrides): + record = {"age": {2027: 63}, "is_tax_unit_head": True} + for component in COMPONENTS: + prefix = f"social_security_{component}" + record[f"{prefix}_reported"] = {2027: 12_000} + record[f"{prefix}_worker_eligibility_year"] = eligibility + record[f"{prefix}_first_cola_year"] = 2027 + record.update(overrides) + return record + + +def _situation(records=None, incomes=None, statuses=None): + records = records or {"worker": _record()} + incomes = incomes or {} + statuses = statuses or {} + situation = { + "people": deepcopy(records), + "tax_units": {}, + "families": {}, + "spm_units": {}, + "households": {}, + } + for person in records: + unit = {"members": [person]} + if person in incomes: + income = incomes[person] + unit["social_security_cola_means_test_magi"] = ( + income if isinstance(income, dict) else {2026: income} + ) + if person in statuses: + status = statuses[person] + unit["filing_status"] = ( + status if isinstance(status, dict) else {2026: status} + ) + unit["tax_unit_itemizes"] = {year: False for year in range(2025, 2031)} + situation["tax_units"][person] = unit + for entity in ("families", "spm_units", "households"): + situation[entity][person] = {"members": [person]} + return situation + + +def _simulation(situation, reform=None, system=None): + return Simulation( + situation=deepcopy(situation), + tax_benefit_system=(system or Simulation.default_tax_benefit_system_instance), + **({"reform": reform} if reform is not None else {}), + ) + + +@pytest.fixture(scope="module") +def means_test_system(): + return _simulation(_situation(), _reform()).tax_benefit_system + + +@pytest.fixture(scope="module") +def disabled_means_test_system(): + return _simulation( + _situation(), + _reform( + enabled=False, + scope="all_beneficiaries", + below_threshold_cola_index="chained_cpi_u", + ), + ).tax_benefit_system + + +def _event(year, rate): + return { + "payment_year": year, + "payment_month": 1, + "rate": rate, + "rounding_rule": 2, + } + + +def test_neutral_defaults_require_explicit_threshold_package(): + p = Simulation.default_tax_benefit_system_instance.parameters("2027-01-01") + means_test = p.gov.contrib.ssa.benefit_levers.cola.means_test + assert not means_test.enabled + assert means_test.scope == 0 + assert means_test.first_eligible_year == 2026 + assert means_test.below_threshold_cola_index == 0 + assert means_test.pre_receipt_cola_index == -1 + assert means_test.thresholds.single == means_test.thresholds.joint == 0 + + +@pytest.mark.parametrize( + "amount,base,current,rounding,expected", + [ + (85_000, 100, 120, 100, 102_000), + (170_000, 100, 120, 100, 204_000), + (85_000, 100, 138, 100, 117_300), + (170_000, 100, 138, 100, 234_600), + (100, 100, 100.5, 1, 101), + (100_000, 100, 100.05, 100, 100_100), + ], +) +def test_cpi_u_threshold_indexing_and_decimal_half_up( + amount, base, current, rounding, expected +): + assert indexed_threshold(amount, base, current, rounding) == expected + + +def _chained_parameters(levels): + return ParameterNode( + data={ + "gov": { + "bls": { + "cpi": { + "c_cpi_u": { + "values": { + f"{year}-07-01": value for year, value in levels.items() + } + } + } + } + } + } + ) + + +def test_chained_index_uses_q3_high_water_and_half_up(): + parameters = _chained_parameters({2000: 100, 2001: 99, 2002: 100.04, 2003: 101}) + rates = chained_cola_rates(parameters, 2004) + assert [rates[year] for year in (2002, 2003, 2004)] == [0, 0, 0.01] + half_up = chained_cola_rates(_chained_parameters({2000: 100, 2001: 100.05}), 2002) + assert half_up[2002] == 0.001 + + +def test_chained_index_reads_all_three_q3_months(): + parameters = _chained_parameters({2000: 100, 2001: 101}) + index = parameters.gov.bls.cpi.c_cpi_u + index.update(start="2001-08-01", value=102) + index.update(start="2001-09-01", value=103) + # Q3 means100 and102 give2%; September alone would give3%. + assert chained_cola_rates(parameters, 2002)[2002] == 0.02 + + +@pytest.mark.parametrize("index", ["statutory", "chained_cpi_u"]) +@pytest.mark.parametrize("scope", ["first_eligible_cohorts", "all_beneficiaries"]) +def test_each_index_scope_combination_and_first_eligible_boundary(index, scope): + records = { + "old_low": _record(2025), + "old_high": _record(2025), + "new_low": _record(2026), + "new_high": _record(2026), + } + incomes = {name: (200_000 if name.endswith("high") else 0) for name in records} + simulation = _simulation( + _situation(records, incomes), + _reform(below_threshold_cola_index=index, scope=scope), + ) + parameters = simulation.tax_benefit_system.parameters + baseline_rate = next( + event["rate"] + for event in cola_events(parameters, 2027) + if event["payment_year"] == 2027 + ) + below_rate = ( + baseline_rate + if index == "statutory" + else chained_cola_rates(parameters, 2027)[2027] + ) + below_factor = (1 + below_rate) / (1 + baseline_rate) + withheld_factor = 1 / (1 + baseline_rate) + expected_withholding = ( + [False, False, False, True] + if scope == "first_eligible_cohorts" + else [False, True, False, True] + ) + expected_factors = ( + [1, 1, below_factor, withheld_factor] + if scope == "first_eligible_cohorts" + else [below_factor, withheld_factor, below_factor, withheld_factor] + ) + for component in COMPONENTS: + np.testing.assert_array_equal( + simulation.calculate(f"social_security_{component}_cola_withheld", 2027), + expected_withholding, + ) + np.testing.assert_allclose( + simulation.calculate( + f"social_security_{component}_joint_cola_factor", 2027 + ), + expected_factors, + ) + + +def test_equality_is_paid_and_each_filing_status_has_an_explicit_group( + means_test_system, +): + statuses = ["SINGLE", "HEAD_OF_HOUSEHOLD", "SEPARATE", "JOINT", "SURVIVING_SPOUSE"] + records = {name: _record() for name in statuses} + records["equality"] = _record() + simulation = _simulation( + _situation( + records, + {**{name: 100_001 for name in statuses}, "equality": 100_000}, + {**{name: name for name in statuses}, "equality": "SINGLE"}, + ), + system=means_test_system, + ) + np.testing.assert_array_equal( + simulation.calculate("social_security_retirement_cola_withheld", 2027), + [True, True, True, False, False, False], + ) + + +def test_surviving_spouse_mapping_is_a_parameter(): + simulation = _simulation( + _situation( + incomes={"worker": 100_001}, statuses={"worker": "SURVIVING_SPOUSE"} + ), + _reform(**{"filing_groups.SURVIVING_SPOUSE": "single"}), + ) + assert simulation.calculate("social_security_retirement_cola_withheld", 2027)[0] + + +def test_only_prior_tax_year_income_and_filing_status_make_the_decision( + means_test_system, +): + simulation = _simulation( + _situation( + incomes={"worker": {2026: 150_000, 2027: 0}}, + statuses={"worker": {2026: "SINGLE", 2027: "JOINT"}}, + ), + system=means_test_system, + ) + assert simulation.calculate("social_security_retirement_cola_withheld", 2027)[0] + + +def test_magi_excludes_social_security_and_retains_irmaa_addbacks(): + situation = _situation() + situation["people"]["worker"].update( + tax_exempt_interest_income={2026: 2_000}, + us_bonds_for_higher_ed={2026: 300}, + ) + situation["people"]["spouse"] = { + "age": {2026: 63}, + "is_tax_unit_spouse": True, + "tax_exempt_interest_income": {2026: 500}, + } + situation["people"]["dependent"] = { + "age": {2026: 19}, + "is_tax_unit_dependent": True, + "tax_exempt_interest_income": {2026: 100_000}, + # The parent's education-bond exclusion can be recorded on the child. + "us_bonds_for_higher_ed": {2026: 700}, + } + members = list(situation["people"]) + unit = situation["tax_units"]["worker"] + unit.update( + members=members, + adjusted_gross_income={2026: 100_000}, + tax_unit_taxable_social_security={2026: 10_000}, + section_911_excluded_income={2026: 4_000}, + specified_possession_income={2026: 5_000}, + puerto_rico_income={2026: 6_000}, + ) + for entity in ("families", "spm_units", "households"): + situation[entity]["worker"]["members"] = members + situation["marital_units"] = { + "couple": {"members": ["worker", "spouse"]}, + "dependent": {"members": ["dependent"]}, + } + simulation = _simulation(situation) + # 100000-10000+2500+1000+4000+5000+6000; the dependent's own interest + # is on a separate return and never enters the parents' income measure. + assert ( + simulation.calculate("social_security_cola_means_test_magi", 2026)[0] == 108_500 + ) + + +def test_supplied_income_history_overrides_the_extended_record_proxy( + means_test_system, +): + records = { + "supplied_history": _record(employment_income={2026: 200_000}), + "extended_record": _record(employment_income={2026: 200_000}), + } + simulation = _simulation( + _situation(records, {"supplied_history": {2026: 90_000}}), + system=means_test_system, + ) + np.testing.assert_array_equal( + simulation.calculate("social_security_retirement_cola_withheld", 2027), + [False, True], + ) + + +@pytest.mark.parametrize( + "sensitivity,income,expected", + [("never_withheld", 1_000_000, False), ("always_withheld", 0, True)], +) +def test_static_cross_section_proxy_sensitivities(sensitivity, income, expected): + simulation = _simulation( + _situation(incomes={"worker": income}), + _reform(income_sensitivity=sensitivity), + ) + assert ( + bool(simulation.calculate("social_security_retirement_cola_withheld", 2027)[0]) + is expected + ) + + +def test_explicit_withholding_input_remains_a_final_decision(means_test_system): + simulation = _simulation( + _situation( + {"worker": _record(social_security_retirement_cola_withheld={2027: True})}, + {"worker": 0}, + ), + system=means_test_system, + ) + assert simulation.calculate("social_security_retirement_cola_withheld", 2027)[0] + assert simulation.calculate("social_security_retirement", 2027)[0] < 12_000 + + +def test_zero_cola_and_downward_adjustment_floor(): + assert cola_event_multiplier(0, 0.0015, True) == 1 + assert cola_event_multiplier(0, -0.003, False) == 1 + assert cola_event_multiplier(0.002, -0.003, False) == pytest.approx(1 / 1.002) + assert cola_event_multiplier(0.002, -0.003, True) == pytest.approx(1 / 1.002) + + +def test_two_year_cpi_e_and_withholding_use_one_joint_numerator(): + events = [_event(2029, 0.03), _event(2030, 0.04)] + factor = joint_cola_factor( + events, (2029, 1), 2030, {2029: 0.0015, 2030: 0.0015}, {2029: True} + ) + # Baseline1000*1.03*1.04=1071.20. First COLA is withheld, so the + # reformed path is1000*1*1.0415=1041.50, rather than retaining CPI-E. + assert factor == pytest.approx(1041.50 / 1071.20) + assert 1071.20 * factor == pytest.approx(1041.50) + separate_factors = ( + joint_cola_factor(events, (2029, 1), 2030, {2029: 0.0015, 2030: 0.0015}, {}) + / 1.03 + ) + assert factor != pytest.approx(separate_factors) + + +@pytest.mark.parametrize( + "below,pre_receipt,use_chained_before_receipt", + [ + ("chained_cpi_u", "follow_below_threshold", True), + ("chained_cpi_u", "statutory", False), + ("statutory", "chained_cpi_u", True), + ], +) +def test_pre_receipt_index_is_separate_and_has_no_income_withholding( + below, pre_receipt, use_chained_before_receipt +): + simulation = _simulation( + _situation( + { + "worker": _record( + social_security_retirement_entitlement_year=2029, + social_security_retirement_entitlement_month=12, + social_security_retirement_reported={2030: 12_000}, + ) + }, + {"worker": {year: 1_000_000 for year in range(2026, 2030)}}, + ), + _reform( + below_threshold_cola_index=below, + pre_receipt_cola_index=pre_receipt, + ), + ) + parameters = simulation.tax_benefit_system.parameters + chained = chained_cola_rates(parameters, 2030) + baseline = { + event["payment_year"]: event["rate"] for event in cola_events(parameters, 2030) + } + expected = prod( + (1 + chained[year]) / (1 + baseline[year]) if use_chained_before_receipt else 1 + for year in (2027, 2028, 2029) + ) / (1 + baseline[2030]) + assert not simulation.calculate("social_security_retirement_cola_withheld", 2029)[0] + assert simulation.calculate("social_security_retirement_cola_withheld", 2030)[0] + assert simulation.calculate("social_security_retirement_joint_cola_factor", 2030)[ + 0 + ] == pytest.approx(expected) + + +def _calculator_record(name): + record = _record() + record.update( + social_security_retirement_source="supplied_pia", + ss_calculator_source="pia", + ss_calculator_worker_id=name, + ss_calculator_birth_date="1964-01-02", + ss_calculator_entitlement_month="2026-01", + ss_calculator_eligibility_year=2026, + ss_calculator_supplied_pia=1000, + ss_calculator_claiming_factor="1", + ss_calculator_insured=True, + ss_calculator_no_special_monthly_test=True, + ss_calculator_no_prior_withholding_adjustment=True, + ss_calculator_earnings_test_no_trigger=True, + ) + return record + + +def test_calculated_source_applies_withholding_and_cpi_e_once(): + reform = _reform() + reform.update( + { + "gov.contrib.ssa.benefit_levers.cola.start_year": {"2020-01-01": 2027}, + "gov.contrib.ssa.benefit_levers.cola.pp_change": {"2020-01-01": 0.0015}, + } + ) + simulation = _simulation( + _situation( + {name: _calculator_record(name) for name in ("paid", "withheld")}, + {"paid": {2026: 0}, "withheld": {2026: 1_000_000}}, + ), + reform, + ) + # Baseline3.6%+.15pp produces PIA1037.5 ->1037 checks. Above the + # line, the single joint numerator produces PIA1000 ->1000 checks. + np.testing.assert_array_equal( + simulation.calculate("ss_calculated_payment", "2027-01"), [1037, 1000] + ) + np.testing.assert_array_equal( + simulation.calculate("social_security_retirement_joint_cola_factor", 2027), + [1, 1], + ) + np.testing.assert_array_equal( + simulation.calculate("social_security_retirement", 2027), + simulation.calculate("ss_calculated_retirement", 2027), + ) + np.testing.assert_array_equal( + simulation.calculate("social_security_retirement", 2027), [12_444, 12_000] + ) + + +def test_midyear_first_check_retains_the_pre_receipt_index(): + records = { + "reported": _record(), + "calculated": _calculator_record("calculated"), + } + for record in records.values(): + record.update( + social_security_retirement_entitlement_year=2028, + social_security_retirement_entitlement_month=6, + ) + records["calculated"]["ss_calculator_entitlement_month"] = "2028-06" + simulation = _simulation( + _situation( + records, + {name: {year: 1_000_000 for year in range(2026, 2029)} for name in records}, + ), + _reform(pre_receipt_cola_index="chained_cpi_u"), + ) + # First receipt is July2028. January2028's COLA therefore uses the + # pre-receipt index regardless of income; January2029 is income-tested. + np.testing.assert_array_equal( + simulation.calculate("social_security_retirement_cola_withheld", 2028), + [False, False], + ) + np.testing.assert_array_equal( + simulation.calculate("social_security_retirement_cola_withheld", 2029), + [True, True], + ) + parameters = simulation.tax_benefit_system.parameters + chained = chained_cola_rates(parameters, 2028) + baseline = { + event["payment_year"]: event["rate"] for event in cola_events(parameters, 2028) + } + expected = prod((1 + chained[year]) / (1 + baseline[year]) for year in (2027, 2028)) + assert simulation.calculate("social_security_retirement_joint_cola_factor", 2028)[ + 0 + ] == pytest.approx(expected) + pia = Decimal(1000) + for year in (2027, 2028): + pia = (pia * (1 + Decimal(str(chained[year])))).quantize( + Decimal("0.1"), rounding=ROUND_DOWN + ) + assert simulation.calculate("ss_calculated_payment", "2028-07")[1] == int(pia) + + +def test_named_withholding_uses_selected_calculator_cohort_and_receipt( + means_test_system, disabled_means_test_system +): + records = {} + for name, birth, eligibility in ( + ("old", "1955-01-02", 2017), + ("new", "1964-01-02", 2026), + ): + # Only calculator dates identify these records; ordinary population + # age, worker-cohort and entitlement inputs are deliberately absent. + records[name] = { + field: value + for field, value in _calculator_record(name).items() + if field.startswith("ss_calculator_") + or field in ("social_security_retirement_source", "is_tax_unit_head") + } + records[name].update( + ss_calculator_birth_date=birth, + ss_calculator_eligibility_year=eligibility, + ss_calculator_entitlement_month="2028-06", + ) + situation = _situation( + records, + {name: {year: 1_000_000 for year in range(2026, 2029)} for name in records}, + ) + simulation = _simulation(situation, system=means_test_system) + baseline = _simulation(situation, system=disabled_means_test_system) + np.testing.assert_array_equal( + simulation.calculate("social_security_retirement_cola_withheld", 2028), + [False, False], + ) + np.testing.assert_array_equal( + simulation.calculate("social_security_retirement_cola_withheld", 2029), + [False, True], + ) + np.testing.assert_array_equal( + simulation.calculate("ss_calculated_payment", "2028-07"), + baseline.calculate("ss_calculated_payment", "2028-07"), + ) + payments = simulation.calculate("ss_calculated_payment", "2029-01") + baseline_payments = baseline.calculate("ss_calculated_payment", "2029-01") + assert payments[0] == baseline_payments[0] + assert payments[1] < baseline_payments[1] + + +def test_a9_preset_reproduces_captured_description(): + simulation = _simulation( + _situation( + { + name: _record( + 2025, + age={2028: 70, 2029: 70}, + social_security_retirement_first_cola_year=2029, + ) + for name in ("below", "above") + }, + {"below": {2027: 0, 2028: 0}, "above": {2027: 1_000_000, 2028: 1_000_000}}, + ), + a9, + ) + p = simulation.tax_benefit_system.parameters("2029-01-01") + means_test = p.gov.contrib.ssa.benefit_levers.cola.means_test + assert means_test.enabled + assert means_test.scope == 1 + assert means_test.effective_year == 2028 + assert means_test.effective_month == 12 + assert means_test.below_threshold_cola_index == 1 + assert means_test.pre_receipt_cola_index == 1 + assert means_test.thresholds.base_year == 2018 + assert means_test.thresholds.single == 85_000 + assert means_test.thresholds.joint == 170_000 + np.testing.assert_array_equal( + simulation.calculate("social_security_retirement_cola_withheld", 2028), + [False, False], + ) + np.testing.assert_array_equal( + simulation.calculate("social_security_retirement_cola_withheld", 2029), + [False, True], + ) + parameters = simulation.tax_benefit_system.parameters + baseline = next( + event["rate"] + for event in cola_events(parameters, 2029) + if event["payment_year"] == 2029 + ) + chained = chained_cola_rates(parameters, 2029)[2029] + np.testing.assert_allclose( + simulation.calculate("social_security_retirement_joint_cola_factor", 2029), + [(1 + chained) / (1 + baseline), 1 / (1 + baseline)], + ) + + +def test_disabled_later_policy_retains_an_earlier_withheld_event(): + reform = _reform() + reform[f"{MEANS_TEST}.enabled"]["2028-01-01"] = False + simulation = _simulation( + _situation(incomes={"worker": {2026: 1_000_000, 2027: 1_000_000}}), reform + ) + baseline = next( + event["rate"] + for event in cola_events(simulation.tax_benefit_system.parameters, 2027) + if event["payment_year"] == 2027 + ) + assert not simulation.calculate("social_security_retirement_cola_withheld", 2028)[0] + assert simulation.calculate("social_security_retirement_joint_cola_factor", 2028)[ + 0 + ] == pytest.approx(1 / (1 + baseline)) + + +@pytest.mark.parametrize("calculated", [False, True], ids=["reported", "calculator"]) +def test_multi_year_income_dependency_graph_is_acyclic(means_test_system, calculated): + years = range(2025, 2031) + record = _record( + age={year: year - 1964 for year in years}, + employment_income={year: 120_000 for year in years}, + traditional_ira_contributions={year: 1_000 for year in years}, + social_security_retirement_reported={year: 12_000 for year in years}, + ) + if calculated: + record.update(_calculator_record("worker")) + record["age"] = {year: year - 1964 for year in years} + record["employment_income"] = {year: 10_000 for year in years} + record["taxable_interest_income"] = {year: 120_000 for year in years} + simulation = _simulation(_situation({"worker": record}), system=means_test_system) + simulation.trace = True + for year in (2028, 2029, 2030): + assert np.isfinite( + simulation.calculate("social_security_retirement", year) + ).all() + edges = {} + + def key(node): + return node.name, str(node.period), getattr(node, "branch_name", None) + + def visit(node): + parent = key(node) + edges.setdefault(parent, set()) + for child in node.children: + edges[parent].add(key(child)) + visit(child) + + for tree in simulation.tracer.trees: + visit(tree) + active, complete = set(), set() + + def assert_dag(node): + assert node not in active, f"Dependency cycle at {node}" + if node in complete: + return + active.add(node) + for child in edges[node]: + assert_dag(child) + active.remove(node) + complete.add(node) + + for node in edges: + assert_dag(node) + income_periods = { + int(period[:4]) + for name, period, _ in edges + if name == "social_security_cola_means_test_magi" + } + assert {2027, 2028, 2029} <= income_periods + income_edges = [] + for parent, children in edges.items(): + if parent[0] != "social_security_cola_means_test_withheld": + continue + for child in children: + if child[0] == "social_security_cola_means_test_magi": + income_edges.append((parent, child)) + assert int(child[1][:4]) == int(parent[1][:4]) - 1 + assert income_edges, "The graph must exercise the actual prior-year income read" + + +@settings(max_examples=12, deadline=None) +@given(lower=st.integers(0, 250_000), increase=st.integers(0, 250_000)) +def test_property_withholding_is_monotone_in_income(means_test_system, lower, increase): + simulation = _simulation( + _situation( + {"lower": _record(), "higher": _record()}, + {"lower": lower, "higher": lower + increase}, + ), + system=means_test_system, + ) + decisions = simulation.calculate("social_security_retirement_cola_withheld", 2027) + assert int(decisions[0]) <= int(decisions[1]) + + +@settings(max_examples=60, deadline=None) +@given( + rate=st.integers(0, 100).map(lambda value: value / 1000), + delta=st.integers(-100, 100).map(lambda value: value / 10000), +) +def test_property_a_withheld_year_never_raises_the_benefit(rate, delta): + withheld = float(cola_event_multiplier(rate, delta, True)) + paid = float(cola_event_multiplier(rate, delta, False)) + assert withheld <= 1 + assert withheld <= paid + + +@settings(max_examples=60, deadline=None) +@given( + values=st.lists( + st.tuples(st.integers(0, 100), st.integers(-100, 100), st.booleans()), + min_size=1, + max_size=8, + ) +) +def test_property_joint_factor_equals_product_of_per_year_ratios(values): + events = [_event(2027 + i, c / 1000) for i, (c, _, _) in enumerate(values)] + deltas = {2027 + i: d / 10000 for i, (_, d, _) in enumerate(values)} + withheld = {2027 + i: w for i, (_, _, w) in enumerate(values)} + expected = prod( + (1 if w else 1 + max(0, c / 1000 + d / 10000)) / (1 + c / 1000) + for c, d, w in values + ) + actual = joint_cola_factor(events, (2027, 1), 2034, deltas, withheld) + assert actual == pytest.approx(expected) + + +@settings(max_examples=8, deadline=None) +@given( + reported=st.integers(0, 100_000), + prior_income=st.integers(0, 1_000_000), +) +def test_property_means_test_off_preserves_stacked_base_outputs( + disabled_means_test_system, reported, prior_income +): + record = _record() + for component in COMPONENTS: + record[f"social_security_{component}"] = {2027: reported} + record.pop(f"social_security_{component}_reported") + situation = _situation({"worker": record}, {"worker": prior_income}) + baseline = _simulation(situation) + disabled = _simulation(situation, system=disabled_means_test_system) + for variable in ( + *(f"social_security_{component}" for component in COMPONENTS), + "social_security", + *(f"social_security_{component}_joint_cola_factor" for component in COMPONENTS), + ): + np.testing.assert_array_equal( + disabled.calculate(variable, 2027), baseline.calculate(variable, 2027) + ) + + +def _assert_cached_outputs_equal(reference, actual): + for population in reference.populations.values(): + for name, holder in population._holders.items(): + for period in holder.get_known_periods(): + expected = holder.get_array(period) + if expected is not None: + np.testing.assert_array_equal( + actual.calculate(name, period), + expected, + err_msg=f"{name} at {period}", + ) + + +@pytest.mark.parametrize("calculated", [False, True], ids=["reported", "calculator"]) +def test_means_test_off_preserves_downstream_dependency_closure( + disabled_means_test_system, calculated +): + record = _calculator_record("worker") if calculated else _record() + if not calculated: + # Differential ingress uses legacy public inputs, as M9 requires. + for component in COMPONENTS: + record[f"social_security_{component}"] = {2027: 6_000} + record.pop(f"social_security_{component}_reported") + record.update( + employment_income={2027: 10_000}, + taxable_interest_income={2027: 20_000}, + ) + situation = _situation({"worker": record}) + situation["households"]["worker"]["state_code"] = {2027: "NY"} + baseline = _simulation(situation) + disabled = _simulation(situation, system=disabled_means_test_system) + for name in ( + "social_security", + "taxable_social_security", + "adjusted_gross_income", + "income_tax", + "net_investment_income_tax", + "medicare_part_b_premium", + "household_net_income", + ): + np.testing.assert_array_equal( + disabled.calculate(name, 2027), baseline.calculate(name, 2027), err_msg=name + ) + if calculated: + for name, period in ( + ("ss_calculated_payment", "2027-01"), + ("ss_calculated_annual_entitlement", 2027), + ("ss_calculated_retirement", 2027), + ): + np.testing.assert_array_equal( + disabled.calculate(name, period), baseline.calculate(name, period) + ) + _assert_cached_outputs_equal(baseline, disabled) diff --git a/policyengine_us/tools/social_security/engine.py b/policyengine_us/tools/social_security/engine.py index 8c779c480b9..0fc219e0c50 100644 --- a/policyengine_us/tools/social_security/engine.py +++ b/policyengine_us/tools/social_security/engine.py @@ -17,6 +17,7 @@ ) from .parameters import bend_points_at, formula_path_at from .pia import FormulaPath +from .means_test import component_withholding, indexed_cola_rate def calculator_output(person, period, parameters, output): @@ -168,9 +169,31 @@ def money(name, when=period): lower[earnings_year] = test.exempt_amount_under_fra higher[earnings_year] = test.exempt_amount_year_of_fra changes, absolute_rates, withheld_years, across, payable = {}, {}, [], {}, {} + # December entitlement is paid the following January. Payment outputs + # for y do not need that future event's y-income decision; eagerly + # reading it would depend on the very checks being calculated. An + # entitlement output can read it through the separate payment route. + policy_horizon = year + int( + output + in ( + "annual_entitlement", + "earnings_test_reduction", + "deferred_earnings_test_payment", + ) + or (output == "benefit" and str(period)[5:7] == "12") + ) + component = "survivors" if source == "survivor" else "retirement" + eligibility = np.full(len(sources), benchmark) + receipt = np.full( + len(sources), + int(record.entitlement_month[:4]) + (record.entitlement_month[5:7] == "12"), + ) + receipt_month = np.full( + len(sources), int(record.entitlement_month[5:7]) % 12 + 1 + ) for event in events: payment_year = int(event.payment_month[:4]) - if payment_year < 1983: + if payment_year < 1983 or payment_year > policy_horizon: continue policy = parameters.gov.contrib.ssa.benefit_levers.cola( f"{payment_year}-01-01" @@ -196,22 +219,61 @@ def money(name, when=period): applies = applies and ( not policy.end_year or payment_year <= policy.end_year ) - withheld = bool( - policy.withholding_enabled - and has_input_for_period( - person.simulation, - "ss_calculator_withhold_cola", - str(payment_year), + withholding_enabled = ( + policy.withholding_enabled or policy.means_test.enabled + ) + if withholding_enabled and has_input_for_period( + person.simulation, + "ss_calculator_withhold_cola", + str(payment_year), + ): + withheld = bool(value("withhold_cola", str(payment_year))) + elif withholding_enabled: + withheld = bool( + component_withholding( + person, + payment_year, + parameters, + component, + eligibility_year=eligibility, + receipt_year=receipt, + receipt_month=receipt_month, + use_explicit=True, + )[index] ) - and value("withhold_cola", str(payment_year)) + else: + withheld = False + index_rate = indexed_cola_rate( + person, + period, + parameters, + component, + payment_year, + float(event.rate), + eligibility_year=eligibility, + receipt_year=receipt, + receipt_month=receipt_month, + ) + resolved_index = Decimal( + str(float(np.broadcast_to(index_rate, len(sources))[index])) ) changes[payment_year] = delta if applies else Decimal(0) - if applies and alternative >= 0: + if (applies and alternative >= 0) or resolved_index != event.rate: # An absolute alternative replaces both 1999's as-paid rate # and its corrected rate. Ordinary baseline+delta changes # instead retain that historical distinction in adjusted_events. absolute_rates[event.payment_month] = ( - Decimal(0) if withheld else max(Decimal(0), alternative + delta) + Decimal(0) + if withheld + else max( + Decimal(0), + ( + alternative + if applies and alternative >= 0 + else resolved_index + ) + + (delta if applies else Decimal(0)), + ) ) if withheld: withheld_years.append(payment_year) diff --git a/policyengine_us/tools/social_security/means_test.py b/policyengine_us/tools/social_security/means_test.py new file mode 100644 index 00000000000..621b2642476 --- /dev/null +++ b/policyengine_us/tools/social_security/means_test.py @@ -0,0 +1,399 @@ +"""Shared payment-year decisions for the contributed COLA means test. + +Tax returns are read only at y-1 for a January y COLA. Population records use +their supplied/extended prior-year copy; this module never substitutes current +AGI for a missing history. Cohort proxies are resolved within each score year. +""" + +from datetime import date +from decimal import Decimal, ROUND_HALF_UP + +import numpy as np +from policyengine_core.periods import period as to_period + +from policyengine_us.tools.general import has_input_for_period + + +def indexed_threshold(amount, base_index, current_index, rounding=100): + """Index a base-dollar line and round half up to the configured dollars.""" + amount, base, current, quantum = map( + lambda value: Decimal(str(value)), + (amount, base_index, current_index, rounding), + ) + if any(not value.is_finite() for value in (amount, base, current, quantum)): + raise ValueError("Means-test threshold inputs must be finite") + if amount <= 0 or base <= 0 or current <= 0 or quantum <= 0: + raise ValueError( + "Means-test thresholds, CPI levels and rounding must be positive" + ) + return float( + (amount * current / base / quantum).quantize(Decimal(1), rounding=ROUND_HALF_UP) + * quantum + ) + + +def policy_settings(parameters, payment_year): + return parameters.gov.contrib.ssa.benefit_levers.cola.means_test( + f"{int(payment_year):04d}-01-01" + ) + + +def parameter_choice(value, choices): + """Decode numeric policy choices supported by Core parameter values.""" + try: + return choices[value] + except (KeyError, TypeError) as error: + raise ValueError( + f"Unknown COLA means-test parameter choice: {value}" + ) from error + + +def policy_active(settings, payment_year): + if not settings.enabled: + return False + if not 1 <= settings.effective_month <= 12: + raise ValueError("Means-test effective month must lie in [1, 12]") + # Contemporary COLAs take effect in December and change January checks. + return (int(payment_year) - 1, 12) >= ( + int(settings.effective_year), + int(settings.effective_month), + ) + + +def chained_cola_rates(parameters, through_year): + """C-CPI-U Q3 rates with A5's high-water, half-up and no-decrease rules. + + Monthly observations are averaged over July--September of the effective + year. Projected snapshots hold the model's annual index observations; they + are an explicit projection approximation, not observed monthly forecasts. + The first available complete Q3 (2000) anchors the payment-2001 series. + Withholding never changes this index's high-water mark. + """ + index = parameters.gov.bls.cpi.c_cpi_u + levels = {} + for payment_year in range(2001, int(through_year) + 1): + levels[payment_year] = ( + sum( + ( + Decimal(str(index(f"{payment_year - 1}-{month:02d}-01"))) + for month in (7, 8, 9) + ), + Decimal(0), + ) + / 3 + ) + if any(not level.is_finite() or level <= 0 for level in levels.values()): + raise ValueError("Chained CPI-U levels must be finite and positive") + if not levels: + return {} + anchor = levels[2001] + rates = {} + for year, level in list(levels.items())[1:]: + rate = max(Decimal(0), level / anchor - 1).quantize( + Decimal("0.001"), rounding=ROUND_HALF_UP + ) + rates[year] = float(rate) + if rate > 0: + anchor = level + return rates + + +def tax_unit_withholding(tax_unit, period, parameters): + """Income decision before beneficiary-specific cohort/receipt restrictions.""" + year = period.start.year + p = policy_settings(parameters, year) + if not policy_active(p, year): + return False + sensitivity = parameter_choice( + p.income_sensitivity, + {0: "prior_year", 1: "never_withheld", 2: "always_withheld"}, + ) + if sensitivity == "never_withheld": + return False + if sensitivity == "always_withheld": + return True + previous = period.offset(-1, "year") + income = tax_unit("social_security_cola_means_test_magi", previous) + status = tax_unit("filing_status", previous) + thresholds = p.thresholds + if thresholds.index_lag < 0: + raise ValueError("Means-test CPI-U lag must be nonnegative") + index = parameters.gov.bls.cpi.cpi_u + base = index(f"{int(thresholds.base_year):04d}-12-01") + current = index(f"{year - int(thresholds.index_lag):04d}-12-01") + lines = np.zeros(len(income)) + for filing_status in status.possible_values: + group = parameter_choice( + getattr(p.filing_groups, filing_status.name), {0: "single", 1: "joint"} + ) + line = indexed_threshold( + getattr(thresholds, group), base, current, thresholds.rounding + ) + lines = np.where(status == filing_status, line, lines) + # A9 leaves equality unspecified. Equality receives a COLA by convention. + return income > lines + + +def calculator_component_basis(person, period, component): + """Read selected calculator dates through the calculator's own validation. + + Only the fields used by SuppliedRecord.benchmark are needed here; monetary + inputs and tax calculations are left to the benefit engine. SurvivorBasis + retains its ordinary non-WINDEX, disability and exact-date assertions. + """ + if component not in ("retirement", "survivors"): + return {} + from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, + ) + + from .household import SuppliedRecord, SurvivorBasis, month_number + + sources = person(f"social_security_{component}_source", period) + reported = sources == SocialSecurityBenefitSource.reported + if np.all(reported): + return {} + calculator_sources = person("ss_calculator_source", period) + allowed = ( + ((SocialSecurityBenefitSource.supplied_ratio_survivor, "survivor"),) + if component == "survivors" + else ( + (SocialSecurityBenefitSource.supplied_aime, "aime"), + (SocialSecurityBenefitSource.supplied_pia, "pia"), + ) + ) + results = {} + for index in np.flatnonzero(~reported): + source = str(calculator_sources[index]) + if not any( + sources[index] == selector.index and source == mode + for selector, mode in allowed + ): + raise ValueError( + "Public Social Security source must match the supported ss_calculator_source" + ) + + def value(name): + scalar = person("ss_calculator_" + name, period)[index] + return scalar.item() if hasattr(scalar, "item") else scalar + + def optional_date(name): + text = str(value(name)) + return date.fromisoformat(text) if text else None + + survivor = None + if source == "survivor": + prior_disability = value("prior_disability_eligibility") + if prior_disability not in (0, 1): + raise ValueError( + "ss_calculator_prior_disability_eligibility must explicitly " + "be 0 (no prior eligibility) or 1 (prior disability eligibility)" + ) + same_year = int(value("same_year_death_before_62")) + if same_year not in (-1, 0, 1): + raise ValueError("Same-year death assertion must be -1, 0 or 1") + survivor = SurvivorBasis( + worker_birth_year=int(value("worker_birth_year")), + death_year=int(value("death_year")), + provenance=str(value("basis_provenance")), + worker_birth_date=optional_date("worker_birth_date"), + death_date=optional_date("death_date"), + authoritative_year=int(value("authoritative_year")) or None, + prior_disability_eligibility=bool(prior_disability), + ordinary_non_windex=bool(value("ordinary_non_windex")), + year_only_not_january_1=bool(value("year_only_not_january_1")), + same_year_death_before_62=None if same_year == -1 else bool(same_year), + ) + record = SuppliedRecord( + source=source, + worker_id=str(value("worker_id")), + birth_date=date.fromisoformat(str(value("birth_date"))), + entitlement_month=str(value("entitlement_month")), + insured=bool(value("insured")), + eligibility_year=int(value("eligibility_year")) or None, + source_basis=str(value("source_basis")), + survivor=survivor, + ) + benchmark = record.benchmark() + check_year, check_month = divmod(month_number(record.entitlement_month) + 1, 12) + results[index] = benchmark, check_year, check_month + 1 + return results + + +def component_eligibility_year(person, period, component): + """Resolve supplied worker links before the documented cross-section proxy.""" + calculated = calculator_component_basis(person, period, component) + prefix = f"social_security_{component}" + eligibility = person(f"{prefix}_worker_eligibility_year", period).copy() + links = person(f"{prefix}_worker_id", period) + shared = person("social_security_worker_computation_year", period) + known = {str(identifier): i for i, identifier in enumerate(person.ids)} + for index, link in enumerate(links): + if index in calculated or not link: + continue + if str(link) not in known: + raise ValueError("Social Security worker_id does not identify a person") + other = known[str(link)] + linked = eligibility[other] or shared[other] + if eligibility[index] and linked and eligibility[index] != linked: + raise ValueError("Conflicting supplied and linked worker computation years") + eligibility[index] = eligibility[index] or linked + year = period.start.year + age = person("age", period) + entitlement = person(f"{prefix}_entitlement_year", period) + if component == "retirement": + proxy = year - age + 62 + elif component == "survivors": + proxy = np.where(age >= 60, year - age + 60, 1975) + else: + proxy = np.where(entitlement > 0, entitlement, 1975) + result = np.where(eligibility > 0, eligibility, proxy) + for index, (benchmark, _, _) in calculated.items(): + result[index] = benchmark + return result + + +def component_receipt_year(person, period, component): + prefix = f"social_security_{component}" + entitlement = person(f"{prefix}_entitlement_year", period) + month = person(f"{prefix}_entitlement_month", period) + # Unknown receipt dates use already-in-payment exposure, as A5 does. + result = np.where(entitlement > 0, entitlement + (month == 12), 0) + for index, (_, check_year, _) in calculator_component_basis( + person, period, component + ).items(): + result[index] = check_year + return result + + +def component_receipt_month(person, period, component): + """First check month following the supplied entitlement month.""" + prefix = f"social_security_{component}" + entitlement = person(f"{prefix}_entitlement_year", period) + month = person(f"{prefix}_entitlement_month", period) + if np.any((entitlement > 0) & ((month < 1) | (month > 12))): + raise ValueError("Social Security entitlement months must lie in [1, 12]") + result = np.where(entitlement > 0, month % 12 + 1, 1) + for index, (_, _, check_month) in calculator_component_basis( + person, period, component + ).items(): + result[index] = check_month + return result + + +def before_receipt(payment_year, receipt_year, receipt_month): + """Whether a January COLA check event precedes the first benefit check.""" + return (payment_year < receipt_year) | ( + (payment_year == receipt_year) & (1 < receipt_month) + ) + + +def cohort_mask(settings, eligibility_year): + scope = parameter_choice( + settings.scope, {0: "first_eligible_cohorts", 1: "all_beneficiaries"} + ) + if scope == "all_beneficiaries": + return np.ones_like(eligibility_year, dtype=bool) + return np.asarray(eligibility_year) >= settings.first_eligible_year + + +def component_withholding( + person, + period, + parameters, + component, + eligibility_year=None, + receipt_year=None, + receipt_month=None, + use_explicit=False, +): + """Apply the TaxUnit decision to this beneficiary and payment event. + + The reported-factor caller supplies its score-year eligibility proxy for + every historical event. It does not re-age a static row within one score. + Explicit per-event withholding inputs retain A5's override channel. + """ + period = to_period(period) + year = period.start.year + p = policy_settings(parameters, year) + if use_explicit and has_input_for_period( + person.simulation, f"social_security_{component}_cola_withheld", period + ): + return person(f"social_security_{component}_cola_withheld", period) + if not policy_active(p, year): + return np.zeros(len(person.ids), dtype=bool) + eligibility_year = ( + component_eligibility_year(person, period, component) + if eligibility_year is None + else np.asarray(eligibility_year) + ) + receipt_year = ( + component_receipt_year(person, period, component) + if receipt_year is None + else np.asarray(receipt_year) + ) + receipt_month = ( + component_receipt_month(person, period, component) + if receipt_month is None + else np.asarray(receipt_month) + ) + eligible = cohort_mask(p, eligibility_year) & ~before_receipt( + year, receipt_year, receipt_month + ) + if not np.any(eligible): + return eligible + return eligible & person.tax_unit( + "social_security_cola_means_test_withheld", period + ) + + +def indexed_cola_rate( + person, + reference_period, + parameters, + component, + payment_year, + baseline_rate, + eligibility_year=None, + receipt_year=None, + receipt_month=None, +): + """Select the below-line or pre-receipt index before A5's policy addition.""" + p = policy_settings(parameters, payment_year) + if not policy_active(p, payment_year): + return baseline_rate + eligibility_year = ( + component_eligibility_year(person, reference_period, component) + if eligibility_year is None + else np.asarray(eligibility_year) + ) + receipt_year = ( + component_receipt_year(person, reference_period, component) + if receipt_year is None + else np.asarray(receipt_year) + ) + receipt_month = ( + component_receipt_month(person, reference_period, component) + if receipt_month is None + else np.asarray(receipt_month) + ) + below = parameter_choice( + p.below_threshold_cola_index, {0: "statutory", 1: "chained_cpi_u"} + ) + before = parameter_choice( + p.pre_receipt_cola_index, + {-1: "follow_below_threshold", 0: "statutory", 1: "chained_cpi_u"}, + ) + if before == "follow_below_threshold": + before = below + chained = cohort_mask(p, eligibility_year) & np.where( + before_receipt(payment_year, receipt_year, receipt_month), + before == "chained_cpi_u", + below == "chained_cpi_u", + ) + if not np.any(chained): + return baseline_rate + rates = chained_cola_rates(parameters, payment_year) + if payment_year not in rates: + raise ValueError("Chained CPI-U COLAs require payment year 2002 or later") + return np.where(chained, rates[payment_year], baseline_rate) diff --git a/policyengine_us/tools/social_security_benefit_levers.py b/policyengine_us/tools/social_security_benefit_levers.py index fb86fa14647..c56f7a1efde 100644 --- a/policyengine_us/tools/social_security_benefit_levers.py +++ b/policyengine_us/tools/social_security_benefit_levers.py @@ -72,9 +72,13 @@ def joint_cola_factor(events, first_payment, through_year, pp_changes, withheld) def cola_event_multiplier(rate, delta, withheld): """Shared scalar/vector algebra for an exposed payment event.""" - if not np.isfinite(rate) or rate < 0 or not np.isfinite(delta): + if ( + not np.all(np.isfinite(rate)) + or np.any(np.asarray(rate) < 0) + or not np.all(np.isfinite(delta)) + ): raise ValueError("COLA rates must be nonnegative and changes finite") - numerator = np.where(withheld, 1, 1 + max(0, rate + delta)) + numerator = np.where(withheld, 1, 1 + np.maximum(0, rate + delta)) return numerator / (1 + rate) @@ -123,12 +127,31 @@ def cola_events(parameters, through_year): def selected_reported_benefit(person, period, component): prefix = f"social_security_{component}" source = person(f"{prefix}_source", period) - if np.any(source != SocialSecurityBenefitSource.reported): + reported = source == SocialSecurityBenefitSource.reported + amount = person(f"{prefix}_reported", period) + if np.all(reported): + return amount + if component not in ("retirement", "survivors"): raise NotImplementedError( - "Calculated Social Security source modes are reserved; " - "this implementation supports the reported route only" + "The calculator supports own retirement and restricted survivor records" + ) + calculator_source = person("ss_calculator_source", period) + allowed = ( + ((SocialSecurityBenefitSource.supplied_ratio_survivor, "survivor"),) + if component == "survivors" + else ( + (SocialSecurityBenefitSource.supplied_aime, "aime"), + (SocialSecurityBenefitSource.supplied_pia, "pia"), ) - return person(f"{prefix}_reported", period) + ) + valid = reported.copy() + for selector, calculator_mode in allowed: + valid |= (source == selector) & (calculator_source == calculator_mode) + if not np.all(valid): + raise ValueError( + "Public Social Security source must match the supported ss_calculator_source" + ) + return np.where(reported, amount, person(f"ss_calculated_{component}", period)) def reported_joint_cola_factor(person, period, parameters, component): @@ -151,6 +174,10 @@ def reported_joint_cola_factor(person, period, parameters, component): for name in ("withholding_enabled", "pp_change", "alternative_rate") for entry in getattr(raw_policy, name).values_list ) + active |= any( + entry.instant_str <= f"{year}-12-31" and bool(entry.value) + for entry in raw_policy.means_test.enabled.values_list + ) if not active: return 1 events = cola_events(parameters, year) @@ -185,6 +212,17 @@ def reported_joint_cola_factor(person, period, parameters, component): elif identifier != "": raise ValueError("Social Security worker_id does not identify a person") entitlement = person(f"{prefix}_entitlement_year", period) + from policyengine_us.tools.social_security.means_test import ( + component_eligibility_year, + component_receipt_month, + component_receipt_year, + component_withholding, + indexed_cola_rate, + ) + + means_eligibility = component_eligibility_year(person, period, component) + receipt_year = component_receipt_year(person, period, component) + receipt_month = component_receipt_month(person, period, component) age = person("age", period) if component == "retirement": eligibility = year - age + 62 @@ -242,14 +280,34 @@ def reported_joint_cola_factor(person, period, parameters, component): raise ValueError( "COLA change must be finite; alternative rate must be nonnegative or -1" ) + index_rate = indexed_cola_rate( + person, + period, + parameters, + component, + payment_year, + baseline, + eligibility_year=means_eligibility, + receipt_year=receipt_year, + receipt_month=receipt_month, + ) rate = ( - (alternative if alternative >= 0 else baseline) + policy.pp_change + (alternative if alternative >= 0 else index_rate) + policy.pp_change if applies - else baseline + else index_rate ) withheld = ( - person(f"{prefix}_cola_withheld", payment_year) - if policy.withholding_enabled + component_withholding( + person, + payment_year, + parameters, + component, + eligibility_year=means_eligibility, + receipt_year=receipt_year, + receipt_month=receipt_month, + use_explicit=True, + ) + if policy.withholding_enabled or policy.means_test.enabled else False ) multiplier = cola_event_multiplier(baseline, rate - baseline, withheld) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_cola_means_test_magi.py b/policyengine_us/variables/gov/ssa/ss/social_security_cola_means_test_magi.py new file mode 100644 index 00000000000..47101a45ad6 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_cola_means_test_magi.py @@ -0,0 +1,45 @@ +from policyengine_us.model_api import * + + +class social_security_cola_means_test_magi(Variable): + value_type = float + entity = TaxUnit + label = "Social Security COLA means-test MAGI excluding Social Security" + unit = USD + definition_period = YEAR + documentation = ( + "Income for the contributed COLA means test in this tax year. The " + "OACT A9-linked proposal uses the IRMAA MAGI definition in Social " + "Security Act section 1839(i)(4): AGI with sections 135, 911, 931 " + "and 933 exclusions added back, plus tax-exempt interest. This " + "reform excludes taxable Social Security from that measure. Income " + "and deductions follow the filer's return, leaving out dependents' " + "own amounts; parent-owned savings bond exclusions recorded on a " + "dependent retain the model's filer-amount convention. The COLA " + "decision reads this variable for the prior tax year, so benefit " + "dependencies move strictly backward in time. Supplying this " + "variable for each tax year overrides the calculation and provides " + "an explicit income history. Dataset extension supplies a " + "same-record prior-year proxy in static population cross-sections, " + "which is not an observed longitudinal income history." + ) + reference = ( + "https://www.ssa.gov/OACT/solvency/provisions/cola.html", + "https://www.ssa.gov/OACT/solvency/SJohnson_20161208.pdf#page=8", + "https://www.congress.gov/bill/114th-congress/house-bill/6489/text", + "https://www.ssa.gov/OP_Home/ssact/title18/1839.htm", + ) + + def formula(tax_unit, period, parameters): + agi_with_exclusions = tax_unit( + "agi_plus_section_911_931_933_exclusions", period + ) + p = parameters(period).gov.irs.ald + other_addbacks = tax_unit_non_dep_add( + tax_unit, + period, + ["tax_exempt_interest_income", "us_bonds_for_higher_ed"], + include_dependents=p.filer_amounts_recorded_on_dependents, + ) + taxable_social_security = tax_unit("tax_unit_taxable_social_security", period) + return agi_with_exclusions + other_addbacks - taxable_social_security diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_cola_means_test_withheld.py b/policyengine_us/variables/gov/ssa/ss/social_security_cola_means_test_withheld.py new file mode 100644 index 00000000000..7b3b5cbeec2 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_cola_means_test_withheld.py @@ -0,0 +1,14 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.means_test import tax_unit_withholding + + +class social_security_cola_means_test_withheld(Variable): + value_type = bool + entity = TaxUnit + definition_period = YEAR + label = "Social Security COLA prior-year income test" + documentation = "Prior-tax-year MAGI strictly exceeds the CPI-U-indexed filing-group line. Beneficiary scope is applied separately." + reference = "https://www.ssa.gov/OACT/solvency/provisions/cola.html" + + def formula(tax_unit, period, parameters): + return tax_unit_withholding(tax_unit, period, parameters) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents.py index 17a31f93045..9ebb13ee4ec 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_dependents.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents.py @@ -13,15 +13,17 @@ class social_security_dependents(Variable): documentation = "Selected source with applicable benefit levers; an explicit public input is a final override." def formula(person, period, parameters): - return ( - person("social_security_dependents_base", period) - * where( - person("social_security_dependents_source", period) - == SocialSecurityBenefitSource.reported, - person("social_security_dependents_joint_cola_factor", period) - * person("social_security_dependents_pia_ratio_factor", period), - 1, - ) + base = person("social_security_dependents_base", period) + reported = ( + person("social_security_dependents_source", period) + == SocialSecurityBenefitSource.reported + ) + reported_factor = ( + person("social_security_dependents_joint_cola_factor", period) + * person("social_security_dependents_pia_ratio_factor", period) * person("social_security_across_the_board_factor", period) * person("social_security_dependents_payable_factor", period) ) + # The calculator applies the complete payment policy before its + # earnings test. Its public route must not repeat any of those factors. + return base * where(reported, reported_factor, 1) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_cola_withheld.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_cola_withheld.py index 99af5c2342a..5570c210b36 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_cola_withheld.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_cola_withheld.py @@ -1,4 +1,5 @@ from policyengine_us.model_api import * +from policyengine_us.tools.social_security.means_test import component_withholding class social_security_dependents_cola_withheld(Variable): @@ -7,4 +8,7 @@ class social_security_dependents_cola_withheld(Variable): entity = Person definition_period = YEAR label = "Social Security dependents cola withheld" - documentation = "Whether this payment year COLA is withheld; inactive unless withholding_enabled is set." + documentation = "Automatic prior-year income decision when the COLA means test is enabled, with explicit per-year decision overrides." + + def formula(person, period, parameters): + return component_withholding(person, period, parameters, "dependents") diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability.py index b9987fb0bf1..75e97db9770 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_disability.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability.py @@ -13,15 +13,17 @@ class social_security_disability(Variable): documentation = "Selected source with applicable benefit levers; an explicit public input is a final override." def formula(person, period, parameters): - return ( - person("social_security_disability_base", period) - * where( - person("social_security_disability_source", period) - == SocialSecurityBenefitSource.reported, - person("social_security_disability_joint_cola_factor", period) - * person("social_security_disability_pia_ratio_factor", period), - 1, - ) + base = person("social_security_disability_base", period) + reported = ( + person("social_security_disability_source", period) + == SocialSecurityBenefitSource.reported + ) + reported_factor = ( + person("social_security_disability_joint_cola_factor", period) + * person("social_security_disability_pia_ratio_factor", period) * person("social_security_across_the_board_factor", period) * person("social_security_disability_payable_factor", period) ) + # The calculator applies the complete payment policy before its + # earnings test. Its public route must not repeat any of those factors. + return base * where(reported, reported_factor, 1) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_cola_withheld.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_cola_withheld.py index 96208e14665..65234ed823b 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_disability_cola_withheld.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_cola_withheld.py @@ -1,4 +1,5 @@ from policyengine_us.model_api import * +from policyengine_us.tools.social_security.means_test import component_withholding class social_security_disability_cola_withheld(Variable): @@ -7,4 +8,7 @@ class social_security_disability_cola_withheld(Variable): entity = Person definition_period = YEAR label = "Social Security disability cola withheld" - documentation = "Whether this payment year COLA is withheld; inactive unless withholding_enabled is set." + documentation = "Automatic prior-year income decision when the COLA means test is enabled, with explicit per-year decision overrides." + + def formula(person, period, parameters): + return component_withholding(person, period, parameters, "disability") diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement.py index 9918ec46f7b..da30fd22ca9 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_retirement.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement.py @@ -13,15 +13,17 @@ class social_security_retirement(Variable): documentation = "Selected source with applicable benefit levers; an explicit public input is a final override." def formula(person, period, parameters): - return ( - person("social_security_retirement_base", period) - * where( - person("social_security_retirement_source", period) - == SocialSecurityBenefitSource.reported, - person("social_security_retirement_joint_cola_factor", period) - * person("social_security_retirement_pia_ratio_factor", period), - 1, - ) + base = person("social_security_retirement_base", period) + reported = ( + person("social_security_retirement_source", period) + == SocialSecurityBenefitSource.reported + ) + reported_factor = ( + person("social_security_retirement_joint_cola_factor", period) + * person("social_security_retirement_pia_ratio_factor", period) * person("social_security_across_the_board_factor", period) * person("social_security_retirement_payable_factor", period) ) + # The calculator applies the complete payment policy before its + # earnings test. Its public route must not repeat any of those factors. + return base * where(reported, reported_factor, 1) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_cola_withheld.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_cola_withheld.py index c448d954234..1051300d454 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_cola_withheld.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_cola_withheld.py @@ -1,4 +1,5 @@ from policyengine_us.model_api import * +from policyengine_us.tools.social_security.means_test import component_withholding class social_security_retirement_cola_withheld(Variable): @@ -7,4 +8,7 @@ class social_security_retirement_cola_withheld(Variable): entity = Person definition_period = YEAR label = "Social Security retirement cola withheld" - documentation = "Whether this payment year COLA is withheld; inactive unless withholding_enabled is set." + documentation = "Automatic prior-year income decision when the COLA means test is enabled, with explicit per-year decision overrides." + + def formula(person, period, parameters): + return component_withholding(person, period, parameters, "retirement") diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors.py index 78d1a8505ff..8e55e45c6b5 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_survivors.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors.py @@ -13,15 +13,17 @@ class social_security_survivors(Variable): documentation = "Selected source with applicable benefit levers; an explicit public input is a final override." def formula(person, period, parameters): - return ( - person("social_security_survivors_base", period) - * where( - person("social_security_survivors_source", period) - == SocialSecurityBenefitSource.reported, - person("social_security_survivors_joint_cola_factor", period) - * person("social_security_survivors_pia_ratio_factor", period), - 1, - ) + base = person("social_security_survivors_base", period) + reported = ( + person("social_security_survivors_source", period) + == SocialSecurityBenefitSource.reported + ) + reported_factor = ( + person("social_security_survivors_joint_cola_factor", period) + * person("social_security_survivors_pia_ratio_factor", period) * person("social_security_across_the_board_factor", period) * person("social_security_survivors_payable_factor", period) ) + # The calculator applies the complete payment policy before its + # earnings test. Its public route must not repeat any of those factors. + return base * where(reported, reported_factor, 1) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_cola_withheld.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_cola_withheld.py index 0e151f3d6da..2f0fe18a6d4 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_cola_withheld.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_cola_withheld.py @@ -1,4 +1,5 @@ from policyengine_us.model_api import * +from policyengine_us.tools.social_security.means_test import component_withholding class social_security_survivors_cola_withheld(Variable): @@ -7,4 +8,7 @@ class social_security_survivors_cola_withheld(Variable): entity = Person definition_period = YEAR label = "Social Security survivors cola withheld" - documentation = "Whether this payment year COLA is withheld; inactive unless withholding_enabled is set." + documentation = "Automatic prior-year income decision when the COLA means test is enabled, with explicit per-year decision overrides." + + def formula(person, period, parameters): + return component_withholding(person, period, parameters, "survivors") From ee08e2f380c73338a871092c62892e8dfcb86b4e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 27 Oct 2026 08:10:52 -0700 Subject: [PATCH 66/76] Preserve disabled COLA behavior and complete boundary coverage --- docs/social-security-cola-means-test.md | 11 ++- .../tests/unit/test_ss_cola_means_test.py | 67 ++++++++++++++----- .../tools/social_security/means_test.py | 12 ++++ .../tools/social_security_benefit_levers.py | 8 +-- .../gov/ssa/ss/social_security_dependents.py | 13 ++-- .../gov/ssa/ss/social_security_disability.py | 13 ++-- .../gov/ssa/ss/social_security_retirement.py | 13 ++-- .../gov/ssa/ss/social_security_survivors.py | 13 ++-- 8 files changed, 110 insertions(+), 40 deletions(-) diff --git a/docs/social-security-cola-means-test.md b/docs/social-security-cola-means-test.md index 5c6e19e62cd..f11c638aca4 100644 --- a/docs/social-security-cola-means-test.md +++ b/docs/social-security-cola-means-test.md @@ -61,8 +61,13 @@ reforms must supply the numbers. `means_test.enabled` activates both the income decision and the selected index without requiring another switch. The enclosing `cola.withholding_enabled` remains the A5 switch for externally resolved withholding decisions when the -means test is disabled. Explicit component/year values can override the -calculated decisions. +means test is disabled. With either channel enabled, explicit component +withholding values and `ss_calculator_withhold_cola` are deliberate decision +overrides. The automatic timing, cohort and receipt restrictions apply to +calculated decisions; supplied decisions retain normal override precedence. +When the automatic policy is inactive, manual component inputs retain their +legacy carry-forward behavior. Calculator-specific decisions apply only in +their explicitly supplied year. The `a9` preset in `policyengine_us.reforms.ssa.cola_means_test` enables the policy and selects: @@ -116,6 +121,8 @@ MAGI history can directly set the dedicated variable for the relevant tax years; supplied histories of its income sources and filing status also take precedence over projected records. A tax unit's decision maps to each beneficiary, subject to that component's exposure and cohort scope. +Core stores complete input arrays for each variable and period; an income +history column must supply a value for every tax unit in that array. Although prior-year AGI can depend on prior-year Social Security, every benefit-related lookback moves to an earlier year. The dependency chain stops diff --git a/policyengine_us/tests/unit/test_ss_cola_means_test.py b/policyengine_us/tests/unit/test_ss_cola_means_test.py index 6c24e640e6a..f570cdf933a 100644 --- a/policyengine_us/tests/unit/test_ss_cola_means_test.py +++ b/policyengine_us/tests/unit/test_ss_cola_means_test.py @@ -348,18 +348,25 @@ def test_magi_excludes_social_security_and_retains_irmaa_addbacks(): def test_supplied_income_history_overrides_the_extended_record_proxy( means_test_system, ): - records = { - "supplied_history": _record(employment_income={2026: 200_000}), - "extended_record": _record(employment_income={2026: 200_000}), - } - simulation = _simulation( - _situation(records, {"supplied_history": {2026: 90_000}}), - system=means_test_system, - ) - np.testing.assert_array_equal( - simulation.calculate("social_security_retirement_cola_withheld", 2027), - [False, True], - ) + # Core stores one complete array per input/period. A partially supplied + # TaxUnit input would fill other units with its default zero, so compare + # the supplied-history and calculated-income situations independently. + decisions = [] + for histories in ({"worker": {2026: 90_000}}, {}): + simulation = _simulation( + _situation( + {"worker": _record(employment_income={2026: 200_000})}, histories + ), + system=means_test_system, + ) + decisions.append( + bool( + simulation.calculate("social_security_retirement_cola_withheld", 2027)[ + 0 + ] + ) + ) + assert decisions == [False, True] @pytest.mark.parametrize( @@ -389,6 +396,15 @@ def test_explicit_withholding_input_remains_a_final_decision(means_test_system): assert simulation.calculate("social_security_retirement", 2027)[0] < 12_000 +def test_disabled_means_test_preserves_manual_input_carry_over(): + simulation = _simulation( + _situation( + {"worker": _record(social_security_retirement_cola_withheld={2026: True})} + ) + ) + assert simulation.calculate("social_security_retirement_cola_withheld", 2027)[0] + + def test_zero_cola_and_downward_adjustment_floor(): assert cola_event_multiplier(0, 0.0015, True) == 1 assert cola_event_multiplier(0, -0.003, False) == 1 @@ -474,7 +490,8 @@ def _calculator_record(name): return record -def test_calculated_source_applies_withholding_and_cpi_e_once(): +@pytest.mark.parametrize("other_payment_factors", [False, True]) +def test_calculated_source_applies_withholding_and_cpi_e_once(other_payment_factors): reform = _reform() reform.update( { @@ -482,6 +499,18 @@ def test_calculated_source_applies_withholding_and_cpi_e_once(): "gov.contrib.ssa.benefit_levers.cola.pp_change": {"2020-01-01": 0.0015}, } ) + if other_payment_factors: + reform.update( + { + "gov.contrib.ssa.benefit_levers.across_the_board.start_year": { + "2020-01-01": 2027 + }, + "gov.contrib.ssa.benefit_levers.across_the_board.rate": { + "2020-01-01": 0.08 + }, + "gov.contrib.ssa.benefit_levers.payable.share": {"2020-01-01": 0.8}, + } + ) simulation = _simulation( _situation( {name: _calculator_record(name) for name in ("paid", "withheld")}, @@ -492,7 +521,8 @@ def test_calculated_source_applies_withholding_and_cpi_e_once(): # Baseline3.6%+.15pp produces PIA1037.5 ->1037 checks. Above the # line, the single joint numerator produces PIA1000 ->1000 checks. np.testing.assert_array_equal( - simulation.calculate("ss_calculated_payment", "2027-01"), [1037, 1000] + simulation.calculate("ss_calculated_payment", "2027-01"), + [896, 864] if other_payment_factors else [1037, 1000], ) np.testing.assert_array_equal( simulation.calculate("social_security_retirement_joint_cola_factor", 2027), @@ -503,7 +533,8 @@ def test_calculated_source_applies_withholding_and_cpi_e_once(): simulation.calculate("ss_calculated_retirement", 2027), ) np.testing.assert_array_equal( - simulation.calculate("social_security_retirement", 2027), [12_444, 12_000] + simulation.calculate("social_security_retirement", 2027), + [10_752, 10_368] if other_payment_factors else [12_444, 12_000], ) @@ -678,9 +709,9 @@ def test_multi_year_income_dependency_graph_is_acyclic(means_test_system, calcul simulation = _simulation(_situation({"worker": record}), system=means_test_system) simulation.trace = True for year in (2028, 2029, 2030): - assert np.isfinite( - simulation.calculate("social_security_retirement", year) - ).all() + benefit = simulation.calculate("social_security_retirement", year) + assert np.isfinite(benefit).all() + assert (benefit > 0).all() edges = {} def key(node): diff --git a/policyengine_us/tools/social_security/means_test.py b/policyengine_us/tools/social_security/means_test.py index 621b2642476..15f93f56c47 100644 --- a/policyengine_us/tools/social_security/means_test.py +++ b/policyengine_us/tools/social_security/means_test.py @@ -321,6 +321,18 @@ def component_withholding( ): return person(f"social_security_{component}_cola_withheld", period) if not policy_active(p, year): + # Before E3, this was a pure annual input. Preserve Core's carry-over + # of supplied manual decisions when the automatic policy is inactive. + # Formula caches must never become manual decision sources. + name = f"social_security_{component}_cola_withheld" + earlier = [ + when + for when in person.simulation._get_set_input_periods(name) + if when.start < period.start + and has_input_for_period(person.simulation, name, when) + ] + if earlier: + return person(name, max(earlier, key=lambda when: when.start)) return np.zeros(len(person.ids), dtype=bool) eligibility_year = ( component_eligibility_year(person, period, component) diff --git a/policyengine_us/tools/social_security_benefit_levers.py b/policyengine_us/tools/social_security_benefit_levers.py index c56f7a1efde..3113aaa5b2f 100644 --- a/policyengine_us/tools/social_security_benefit_levers.py +++ b/policyengine_us/tools/social_security_benefit_levers.py @@ -104,7 +104,7 @@ def cola_events(parameters, through_year): last_actual = max(e["payment_year"] for e in events) if through_year > last_actual: indices = { - year: parameters(f"{year}-01-01").gov.ssa.uprating + year: parameters.gov.ssa.uprating(f"{year}-01-01") for year in range(last_actual, through_year + 1) } for year, rate in projected_cola_rates(indices).items(): @@ -162,10 +162,10 @@ def reported_joint_cola_factor(person, period, parameters, component): if not np.any(reported): return 1 year = period.start.year - p = parameters(period).gov.contrib.ssa.benefit_levers.cola + raw_policy = parameters.gov.contrib.ssa.benefit_levers.cola + p = raw_policy(str(period.start)) if not 1 <= p.effective_month <= 12: raise ValueError("COLA effective month must lie in [1, 12]") - raw_policy = parameters.gov.contrib.ssa.benefit_levers.cola # An earlier withheld or adjusted COLA remains in the payment level after # a policy is switched off. Inspect the parameter history, not just today. active = any( @@ -259,7 +259,7 @@ def reported_joint_cola_factor(person, period, parameters, component): if not np.any(exposed): continue validate_cola_rounding_rule(event["rounding_rule"]) - policy = parameters(f"{payment_year}-01-01").gov.contrib.ssa.benefit_levers.cola + policy = raw_policy(f"{payment_year}-01-01") effective = event["effective_date"] applies = ( ( diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents.py index 9ebb13ee4ec..32e84917faf 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_dependents.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents.py @@ -18,12 +18,17 @@ def formula(person, period, parameters): person("social_security_dependents_source", period) == SocialSecurityBenefitSource.reported ) - reported_factor = ( - person("social_security_dependents_joint_cola_factor", period) - * person("social_security_dependents_pia_ratio_factor", period) + reported_benefit = ( + base + * where( + reported, + person("social_security_dependents_joint_cola_factor", period) + * person("social_security_dependents_pia_ratio_factor", period), + 1, + ) * person("social_security_across_the_board_factor", period) * person("social_security_dependents_payable_factor", period) ) # The calculator applies the complete payment policy before its # earnings test. Its public route must not repeat any of those factors. - return base * where(reported, reported_factor, 1) + return where(reported, reported_benefit, base) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability.py index 75e97db9770..d5c1068b80b 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_disability.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability.py @@ -18,12 +18,17 @@ def formula(person, period, parameters): person("social_security_disability_source", period) == SocialSecurityBenefitSource.reported ) - reported_factor = ( - person("social_security_disability_joint_cola_factor", period) - * person("social_security_disability_pia_ratio_factor", period) + reported_benefit = ( + base + * where( + reported, + person("social_security_disability_joint_cola_factor", period) + * person("social_security_disability_pia_ratio_factor", period), + 1, + ) * person("social_security_across_the_board_factor", period) * person("social_security_disability_payable_factor", period) ) # The calculator applies the complete payment policy before its # earnings test. Its public route must not repeat any of those factors. - return base * where(reported, reported_factor, 1) + return where(reported, reported_benefit, base) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement.py index da30fd22ca9..6d3eaf99967 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_retirement.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement.py @@ -18,12 +18,17 @@ def formula(person, period, parameters): person("social_security_retirement_source", period) == SocialSecurityBenefitSource.reported ) - reported_factor = ( - person("social_security_retirement_joint_cola_factor", period) - * person("social_security_retirement_pia_ratio_factor", period) + reported_benefit = ( + base + * where( + reported, + person("social_security_retirement_joint_cola_factor", period) + * person("social_security_retirement_pia_ratio_factor", period), + 1, + ) * person("social_security_across_the_board_factor", period) * person("social_security_retirement_payable_factor", period) ) # The calculator applies the complete payment policy before its # earnings test. Its public route must not repeat any of those factors. - return base * where(reported, reported_factor, 1) + return where(reported, reported_benefit, base) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors.py index 8e55e45c6b5..60883a30b2c 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_survivors.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors.py @@ -18,12 +18,17 @@ def formula(person, period, parameters): person("social_security_survivors_source", period) == SocialSecurityBenefitSource.reported ) - reported_factor = ( - person("social_security_survivors_joint_cola_factor", period) - * person("social_security_survivors_pia_ratio_factor", period) + reported_benefit = ( + base + * where( + reported, + person("social_security_survivors_joint_cola_factor", period) + * person("social_security_survivors_pia_ratio_factor", period), + 1, + ) * person("social_security_across_the_board_factor", period) * person("social_security_survivors_payable_factor", period) ) # The calculator applies the complete payment policy before its # earnings test. Its public route must not repeat any of those factors. - return base * where(reported, reported_factor, 1) + return where(reported, reported_benefit, base) From 0f7917a499397787a0a7b0d9dea16168e6d531a9 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 27 Oct 2026 08:38:37 -0700 Subject: [PATCH 67/76] Test cohort proxies within each static scoring cross-section --- .../tests/unit/test_ss_cola_means_test.py | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/policyengine_us/tests/unit/test_ss_cola_means_test.py b/policyengine_us/tests/unit/test_ss_cola_means_test.py index f570cdf933a..6edd6e278f9 100644 --- a/policyengine_us/tests/unit/test_ss_cola_means_test.py +++ b/policyengine_us/tests/unit/test_ss_cola_means_test.py @@ -290,6 +290,34 @@ def test_surviving_spouse_mapping_is_a_parameter(): assert simulation.calculate("social_security_retirement_cola_withheld", 2027)[0] +def test_static_cohort_proxy_is_reconstructed_within_each_score_year(means_test_system): + record = _record(age={2026: 70, 2035: 70}) + for component in COMPONENTS: + record.pop(f"social_security_{component}_worker_eligibility_year") + record.pop(f"social_security_{component}_first_cola_year") + simulation = _simulation( + _situation( + {"worker": record}, + {"worker": {year: 1_000_000 for year in range(2025, 2035)}}, + ), + system=means_test_system, + ) + # Fixed-age static cross-sections infer E2018 in2026 and E2027 in2035. + # Within the2035 score, all exposed events use E2027 rather than re-aging + # the row separately at each historical payment event. + assert ( + simulation.calculate("social_security_retirement_joint_cola_factor", 2026)[0] + == 1 + ) + events = cola_events(simulation.tax_benefit_system.parameters, 2035) + expected = prod( + 1 / (1 + event["rate"]) for event in events if event["payment_year"] >= 2028 + ) + assert simulation.calculate("social_security_retirement_joint_cola_factor", 2035)[ + 0 + ] == pytest.approx(expected) + + def test_only_prior_tax_year_income_and_filing_status_make_the_decision( means_test_system, ): From ce94d2fe3a3e6ffaae361becc27a924b687d7257 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 27 Oct 2026 09:04:26 -0700 Subject: [PATCH 68/76] Clarify when COLA income thresholds are required --- docs/social-security-cola-means-test.md | 8 ++++++-- .../ssa/benefit_levers/cola/means_test/enabled.yaml | 2 +- .../benefit_levers/cola/means_test/thresholds/joint.yaml | 2 +- .../benefit_levers/cola/means_test/thresholds/single.yaml | 2 +- 4 files changed, 9 insertions(+), 5 deletions(-) diff --git a/docs/social-security-cola-means-test.md b/docs/social-security-cola-means-test.md index f11c638aca4..7e49eeba5a3 100644 --- a/docs/social-security-cola-means-test.md +++ b/docs/social-security-cola-means-test.md @@ -4,7 +4,7 @@ The contributed COLA means test is disabled under current law. It can withhold an annual COLA based on the beneficiary's tax unit, while retaining the benefit level reached before that event. It supports both reported benefit components and the supplied-record household calculator. Scenario thresholds and dates -must be supplied before scoring; PE-US does not choose package values. +come from the scoring settings; PE-US does not choose package values. ## Policy basis and choices @@ -49,7 +49,7 @@ reforms must supply the numbers. | `pre_receipt_cola_index` | `-1` = `follow_below_threshold` | Follow the below-line choice, or explicitly select `0` = `statutory` or `1` = `chained_cpi_u`. | | `income_sensitivity` | `0` = `prior_year` | Alternatively, `1` = `never_withheld` or `2` = `always_withheld`. | | `thresholds.base_year` | `2026` | Dollar year of the supplied lines. | -| `thresholds.single`, `thresholds.joint` | `0`, `0` | Neutral placeholders; activation requires positive scenario lines. | +| `thresholds.single`, `thresholds.joint` | `0`, `0` | Neutral placeholders; evaluated `prior_year` income decisions require positive scenario lines. | | `thresholds.index_lag` | `1` | Years before the payment year for the December CPI-U observation. | | `thresholds.rounding` | `100` | Dollar increment, rounded half up. | | `filing_groups.SINGLE` | `0` = `single` | Select the single line; `1` selects joint. | @@ -69,6 +69,10 @@ When the automatic policy is inactive, manual component inputs retain their legacy carry-forward behavior. Calculator-specific decisions apply only in their explicitly supplied year. +Positive single and joint lines are required when a `prior_year` income +decision is evaluated. The `never_withheld` and `always_withheld` sensitivities +and pre-receipt index-only events do not read those lines. + The `a9` preset in `policyengine_us.reforms.ssa.cola_means_test` enables the policy and selects: diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/enabled.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/enabled.yaml index f511d99f6cf..e7e51312b3e 100644 --- a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/enabled.yaml +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/enabled.yaml @@ -1,4 +1,4 @@ -description: Enable the contributed Social Security COLA means test. Activation requires positive single and joint income thresholds supplied by the scenario. +description: Enable the contributed Social Security COLA means test. Positive single and joint income thresholds supplied by the scenario are required when a prior_year income decision is evaluated. The never_withheld and always_withheld sensitivities and pre-receipt index-only events do not read the thresholds. values: 0000-01-01: false metadata: diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/joint.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/joint.yaml index af7e663bc01..b021f64445a 100644 --- a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/joint.yaml +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/joint.yaml @@ -1,4 +1,4 @@ -description: Joint-group COLA means-test income threshold in base-year dollars. Zero is an unset sentinel; activation requires a positive scenario value. The documented A9 example is 170000 in 2018 dollars, not a default package value. +description: Joint-group COLA means-test income threshold in base-year dollars. Zero is an unset sentinel; a positive scenario value is required when a prior_year income decision is evaluated. The never_withheld and always_withheld sensitivities and pre-receipt index-only events do not read this threshold. The documented A9 example is 170000 in 2018 dollars, not a default package value. values: 0000-01-01: 0 metadata: diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/single.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/single.yaml index 73bd9c7d326..8bdbd07db93 100644 --- a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/single.yaml +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/thresholds/single.yaml @@ -1,4 +1,4 @@ -description: Single-group COLA means-test income threshold in base-year dollars. Zero is an unset sentinel; activation requires a positive scenario value. The documented A9 example is 85000 in 2018 dollars, not a default package value. +description: Single-group COLA means-test income threshold in base-year dollars. Zero is an unset sentinel; a positive scenario value is required when a prior_year income decision is evaluated. The never_withheld and always_withheld sensitivities and pre-receipt index-only events do not read this threshold. The documented A9 example is 85000 in 2018 dollars, not a default package value. values: 0000-01-01: 0 metadata: From 9d9cf5c815c289bbbc07e93947814907076f7b64 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 27 Oct 2026 11:39:37 -0700 Subject: [PATCH 69/76] Bind the A9 preset modifier for shared systems --- policyengine_us/reforms/ssa/cola_means_test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/policyengine_us/reforms/ssa/cola_means_test.py b/policyengine_us/reforms/ssa/cola_means_test.py index 8252a4b012a..348b885fb49 100644 --- a/policyengine_us/reforms/ssa/cola_means_test.py +++ b/policyengine_us/reforms/ssa/cola_means_test.py @@ -5,7 +5,8 @@ class a9(Reform): def apply(self): - self.modify_parameters(self.configure_parameters) + # Shared systems call this unbound method with the system as self. + self.modify_parameters(a9.configure_parameters) @staticmethod def configure_parameters(parameters): From 4f336b78e2cc687b46f83d16e0067d7108df82eb Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 29 Oct 2026 07:59:42 -0400 Subject: [PATCH 70/76] Align calculated benefit fixture with complete payment base --- .../ssa/benefit_levers/social_security_retirement.yaml | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement.yaml index d659a6ea1ef..38836902520 100644 --- a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement.yaml +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement.yaml @@ -99,7 +99,7 @@ output: social_security_dependents: 6000 -- name: Calculated-base extension point bypasses population PIA and COLA factors +- name: Complete calculated base bypasses every population benefit factor period: 2027 input: age: 70 @@ -109,13 +109,15 @@ gov.contrib.ssa.benefit_levers.across_the_board.rate: 0.1 gov.contrib.ssa.benefit_levers.payable.share: 0.8 social_security_retirement_source: supplied_pia - # The route itself is reserved. A supplied selected base stands in for - # the future service's amount already containing its reformed PIA/COLAs. + # This replaces the calculator's complete checks, including its PIA, + # COLAs, across-the-board change and payable share before the earnings test. social_security_retirement_base: 10000 social_security_retirement_pia_ratio_factor: 1.2 social_security_retirement_joint_cola_factor: 1.3 output: - social_security_retirement: 8800 + # The public route returns 10000 * 1. Applying 1.1 * 0.8 again would + # incorrectly reduce these already calculated checks to 8800. + social_security_retirement: 10000 - name: Selecting an OASI-only basis before depletion remains neutral period: 2027 From 2c142b89d48ead3c29786a150b0ec570e1d65e08 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 29 Oct 2026 08:39:03 -0400 Subject: [PATCH 71/76] Preserve legacy COLA route when means test is disabled --- .../tests/core/test_social_security_inputs.py | 6 ++- .../test_social_security_lever_properties.py | 1 + .../tools/social_security_benefit_levers.py | 45 ++++++++++++------- 3 files changed, 34 insertions(+), 18 deletions(-) diff --git a/policyengine_us/tests/core/test_social_security_inputs.py b/policyengine_us/tests/core/test_social_security_inputs.py index 3f5328e40d5..5392c0e8b45 100644 --- a/policyengine_us/tests/core/test_social_security_inputs.py +++ b/policyengine_us/tests/core/test_social_security_inputs.py @@ -1130,12 +1130,14 @@ def test_reported_twins_are_pure_and_original_dataset_names_keep_their_uprater() @pytest.mark.parametrize( "source", ["supplied_aime", "supplied_pia", "supplied_ratio_survivor", "history"] ) -def test_reserved_sources_fail_explicitly_instead_of_using_reported_fallback(source): +def test_unmatched_sources_fail_explicitly_instead_of_using_reported_fallback(source): public = SOCIAL_SECURITY_COMPONENTS[0] simulation = make_simulation( inputs={f"{public}_source": {2024: source}, f"{public}_reported": {2024: 0}} ) - with pytest.raises(NotImplementedError, match="reserved"): + with pytest.raises( + ValueError, match="must match the supported ss_calculator_source" + ): simulation.calculate(public, 2024) diff --git a/policyengine_us/tests/core/test_social_security_lever_properties.py b/policyengine_us/tests/core/test_social_security_lever_properties.py index f2488eb7766..e71ed53b335 100644 --- a/policyengine_us/tests/core/test_social_security_lever_properties.py +++ b/policyengine_us/tests/core/test_social_security_lever_properties.py @@ -280,6 +280,7 @@ def exposure_parameters(**overrides): name: {"values": value if isinstance(value, dict) else {"0000-01-01": value}} for name, value in values.items() } + cola["means_test"] = {"enabled": {"values": {"0000-01-01": False}}} params = ParameterNode( data={ "gov": { diff --git a/policyengine_us/tools/social_security_benefit_levers.py b/policyengine_us/tools/social_security_benefit_levers.py index 3113aaa5b2f..6fee8482cf9 100644 --- a/policyengine_us/tools/social_security_benefit_levers.py +++ b/policyengine_us/tools/social_security_benefit_levers.py @@ -174,10 +174,11 @@ def reported_joint_cola_factor(person, period, parameters, component): for name in ("withholding_enabled", "pp_change", "alternative_rate") for entry in getattr(raw_policy, name).values_list ) - active |= any( + means_test_ever_enabled = any( entry.instant_str <= f"{year}-12-31" and bool(entry.value) for entry in raw_policy.means_test.enabled.values_list ) + active |= means_test_ever_enabled if not active: return 1 events = cola_events(parameters, year) @@ -220,9 +221,13 @@ def reported_joint_cola_factor(person, period, parameters, component): indexed_cola_rate, ) - means_eligibility = component_eligibility_year(person, period, component) - receipt_year = component_receipt_year(person, period, component) - receipt_month = component_receipt_month(person, period, component) + # A5's manual channel needs no calculator or tax-history inputs. Keep + # calculated rows outside population validation when E3 has never applied. + means_eligibility = receipt_year = receipt_month = None + if means_test_ever_enabled: + means_eligibility = component_eligibility_year(person, period, component) + receipt_year = component_receipt_year(person, period, component) + receipt_month = component_receipt_month(person, period, component) age = person("age", period) if component == "retirement": eligibility = year - age + 62 @@ -280,16 +285,20 @@ def reported_joint_cola_factor(person, period, parameters, component): raise ValueError( "COLA change must be finite; alternative rate must be nonnegative or -1" ) - index_rate = indexed_cola_rate( - person, - period, - parameters, - component, - payment_year, - baseline, - eligibility_year=means_eligibility, - receipt_year=receipt_year, - receipt_month=receipt_month, + index_rate = ( + indexed_cola_rate( + person, + period, + parameters, + component, + payment_year, + baseline, + eligibility_year=means_eligibility, + receipt_year=receipt_year, + receipt_month=receipt_month, + ) + if policy.means_test.enabled + else baseline ) rate = ( (alternative if alternative >= 0 else index_rate) + policy.pp_change @@ -307,8 +316,12 @@ def reported_joint_cola_factor(person, period, parameters, component): receipt_month=receipt_month, use_explicit=True, ) - if policy.withholding_enabled or policy.means_test.enabled - else False + if policy.means_test.enabled + else ( + person(f"{prefix}_cola_withheld", payment_year) + if policy.withholding_enabled + else False + ) ) multiplier = cola_event_multiplier(baseline, rate - baseline, withheld) result *= np.where(exposed, multiplier, 1) From ddd42668a6a0ae1f989696392619c660d13741e2 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 30 Oct 2026 09:54:48 -0400 Subject: [PATCH 72/76] Fix independent COLA decisions and preserve legacy calculator withholding --- changelog.d/ss-cola-means-test-e3.added.md | 2 +- docs/social-security-cola-means-test.md | 82 ++-- .../cola/means_test/first_eligible_year.yaml | 5 +- .../social_security_cola_means_test.yaml | 2 +- .../tests/unit/test_ss_cola_means_test.py | 352 ++++++++++++++---- .../tools/social_security/engine.py | 36 +- .../tools/social_security/means_test.py | 49 +-- .../ssa/social_security/calculator/inputs.py | 13 +- ...ocial_security_dependents_cola_withheld.py | 21 +- ...ocial_security_disability_cola_withheld.py | 21 +- ...ocial_security_retirement_cola_withheld.py | 21 +- ...social_security_survivors_cola_withheld.py | 21 +- 12 files changed, 479 insertions(+), 146 deletions(-) diff --git a/changelog.d/ss-cola-means-test-e3.added.md b/changelog.d/ss-cola-means-test-e3.added.md index 9b79cb1d69d..0f034bc7f21 100644 --- a/changelog.d/ss-cola-means-test-e3.added.md +++ b/changelog.d/ss-cola-means-test-e3.added.md @@ -1 +1 @@ -Add an opt-in Social Security COLA means test using prior-year MAGI excluding Social Security, configurable indexed thresholds and filing groups, cohort or all-beneficiary scope, statutory or chained CPI-U COLAs, income-history sensitivities, and an A9 preset. Apply each withholding decision jointly with other COLA changes in reported benefits and the household calculator, with current-law defaults unchanged. +Add an opt-in Social Security COLA means test using prior-year MAGI excluding Social Security, configurable indexed thresholds and filing groups, cohort or all-beneficiary scope, statutory or chained CPI-U COLAs, income-history sensitivities, and an A9 preset. Apply each withholding decision jointly with other COLA changes in reported benefits and the household calculator, with independent per-person overrides while preserving current-law defaults and the legacy A5 withholding input channels. Document the latest-authored-value CPI-U lookup used to index projected thresholds. diff --git a/docs/social-security-cola-means-test.md b/docs/social-security-cola-means-test.md index 7e49eeba5a3..0a28c2eb602 100644 --- a/docs/social-security-cola-means-test.md +++ b/docs/social-security-cola-means-test.md @@ -18,9 +18,10 @@ all beneficiaries beginning with the December 2028 COLA. Its single, head-of-household and separate filers share a line; joint filers have twice that line. It also specifies chained CPI-U before benefit receipt. Its example lines are $85,000 and $170,000 in 2018 dollars, indexed by CPI-U. The preset -below implements these mechanics. OACT's displayed December 2028 dollar -estimates use its projections; this model calculates indexed lines from its -own pinned CPI-U path. +below implements these mechanics. Its 2029 income lines are $117,800 single +and $235,700 joint. OACT's displayed December 2028 lines of $117,300 and +$234,600 use a different CPI projection; this model uses the authored CPI-U +path described below. The default enabled scenario differs from A9 in two parameterized ways: it retains the statutory COLA below the line and restricts the policy to people @@ -50,7 +51,7 @@ reforms must supply the numbers. | `income_sensitivity` | `0` = `prior_year` | Alternatively, `1` = `never_withheld` or `2` = `always_withheld`. | | `thresholds.base_year` | `2026` | Dollar year of the supplied lines. | | `thresholds.single`, `thresholds.joint` | `0`, `0` | Neutral placeholders; evaluated `prior_year` income decisions require positive scenario lines. | -| `thresholds.index_lag` | `1` | Years before the payment year for the December CPI-U observation. | +| `thresholds.index_lag` | `1` | Years before the payment year for the December 1 CPI-U lookup. | | `thresholds.rounding` | `100` | Dollar increment, rounded half up. | | `filing_groups.SINGLE` | `0` = `single` | Select the single line; `1` selects joint. | | `filing_groups.HEAD_OF_HOUSEHOLD` | `0` = `single` | Select the single line; `1` selects joint. | @@ -61,13 +62,21 @@ reforms must supply the numbers. `means_test.enabled` activates both the income decision and the selected index without requiring another switch. The enclosing `cola.withholding_enabled` remains the A5 switch for externally resolved withholding decisions when the -means test is disabled. With either channel enabled, explicit component -withholding values and `ss_calculator_withhold_cola` are deliberate decision -overrides. The automatic timing, cohort and receipt restrictions apply to -calculated decisions; supplied decisions retain normal override precedence. -When the automatic policy is inactive, manual component inputs retain their -legacy carry-forward behavior. Calculator-specific decisions apply only in -their explicitly supplied year. +means test is disabled. With the means test enabled, use the separate annual +Person inputs `social_security_X_cola_override` and +`ss_calculator_cola_override`: `-1` means automatic (the default), `0` means +paid, and `1` means withheld. Each person's default remains automatic when +another person supplies an override. The automatic timing, cohort and receipt +restrictions apply to automatic decisions; an explicit override supplies the +final decision for that person and year. The tri-state overrides expire after +their supplied year; later years return to automatic decisions. + +With the means test disabled, the reported route retains A5's boolean +`social_security_X_cola_withheld` inputs and their legacy carry-forward +behavior. The calculator reads only `ss_calculator_withhold_cola`, in its +explicitly supplied year. Component boolean inputs do not reach the +calculator. These legacy boolean channels do not override an enabled means +test; use the separate tri-state inputs for that purpose. Positive single and joint lines are required when a `prior_year` income decision is evaluated. The `never_withheld` and `always_withheld` sensitivities @@ -161,13 +170,23 @@ line(g, y) = round_half_up_to_R( withheld(y) = prior_year_MAGI > line(g, y) ``` -The CPI-U series is `gov.bls.cpi.cpi_u`. The default observes the December -before the check year, including December 2028 for January 2029 checks. -Parameter snapshots are held until the next authored observation, including -in projections. The December snapshot, one-year payment lag and nearest-$100 -half-up rounding are documented modeling conventions: A9 specifies CPI-U -indexing without prescribing these details. Pin both the index path and -these settings for a score. +The code reads `gov.bls.cpi.cpi_u` at December 1 of year `y - L`, and at +December 1 of the base year. Under the parameter lookup's ordinary semantics, +this returns a BLS December figure where one exists, otherwise the latest +authored value. The series currently contains BLS monthly observations +through June 2026, so `cpi_u("2026-12-01")` returns June 2026's `332.568`. +For later years it uses CBO's February-keyed projection: for example, +`cpi_u("2028-12-01")` returns `348.3`, authored at `2028-02-01`. +These reads do not construct a separate December projection series. + +The default one-year lag uses the December 1, 2028 lookup for January 2029 +checks. The A9 preset's $85,000/$170,000 base lines, December 2018 value +`251.233`, and projected value `348.3` give $117,800/$235,700 after nearest-$100 +half-up rounding. OACT's displayed $117,300/$234,600 lines use a different +CPI projection. The December 1 lookup, one-year payment lag and rounding are +documented modeling conventions: A9 specifies CPI-U indexing without +prescribing these details. Pin both the authored index path and these +settings for a score. The chained option uses Q3 averages of July, August and September levels from `gov.bls.cpi.c_cpi_u`. As with the statutory CPI-W schedule, positive growth @@ -194,9 +213,12 @@ floored at zero. Withholding overrides the entire resulting increase. ## One joint COLA adjustment -Each annual Person variable `social_security_X_cola_withheld`, for `X` equal -to `retirement`, `disability`, `survivors` or `dependents`, has a formula that -resolves its tax unit's decision. Explicit values override that formula. +Each annual Person variable `social_security_X_cola_means_test_withheld`, for +`X` equal to `retirement`, `disability`, `survivors` or `dependents`, resolves +its tax unit's automatic decision and that person's separate +`social_security_X_cola_override` input. The legacy +`social_security_X_cola_withheld` boolean remains A5's input channel when the +means test is disabled. For a reported benefit, the population adjustment accumulates one ratio for each exposed event: @@ -218,8 +240,10 @@ multiplied withholding factor would give the wrong result. The household calculator applies the same per-year decisions inside its COLA schedule. It rounds the resulting PIA under the applicable event rule and -applies claiming and payment calculations once. Its explicit annual -`ss_calculator_withhold_cola` value takes precedence over the shared decision. +applies claiming and payment calculations once. With the means test enabled, +the annual `ss_calculator_cola_override` takes precedence over the component +override, followed by the automatic decision. With the means test disabled, +it retains A5's exact-year `ss_calculator_withhold_cola` channel. For a public retirement component, select `social_security_retirement_source` as `supplied_aime` or `supplied_pia`, and supply the matching @@ -245,10 +269,12 @@ adjustments, high-water years and combined CPI-E changes and withholding. Multi-year graph checks establish that the prior-year income read has no cycle. Calculated-source cases check that the COLA policy is applied once. -Property checks require disabled-policy equality to the stacked base, -nonincreasing benefits when a year is withheld, monotone withholding as income -rises at a fixed line, and equality of the joint factor to the product of its -per-year ratios. Differential checks cover reported situations, downstream -outputs and household calculator examples with this means test disabled. +Compatibility checks pin literal results from main for disabled-policy +reported and calculator routes, including A5 inputs. Other checks require +nonincreasing benefits when a year is withheld and monotone withholding as +income rises at a fixed line. Multi-person cases preserve each person's +automatic decision when someone else supplies an override. Differential +checks cover reported situations, downstream outputs and household +calculator examples with this means test disabled. Population impact requires a separately pinned scoring run; these mechanics do not substitute OACT's solvency estimates for a PE-US result. diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/first_eligible_year.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/first_eligible_year.yaml index 1288983bef2..5f360b0a87e 100644 --- a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/first_eligible_year.yaml +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/cola/means_test/first_eligible_year.yaml @@ -1,7 +1,10 @@ -description: First eligibility cohort reached when the Social Security COLA means-test scope is first_eligible_cohorts. This contributed reform is disabled by default. +description: First eligibility cohort reached when the Social Security COLA means-test scope is first_eligible_cohorts. The 2026 boundary is a configurable scenario choice; the OACT A9 preset applies to all beneficiaries. This contributed reform is disabled by default. values: 0000-01-01: 2026 metadata: unit: year label: Social Security COLA means-test first eligibility cohort economy: false + reference: + - title: OACT provision A9, 2026 Trustees Report basis + href: https://www.ssa.gov/OACT/solvency/provisions/cola.html diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_cola_means_test.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_cola_means_test.yaml index 3cf0c634490..445edaa873c 100644 --- a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_cola_means_test.yaml +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_cola_means_test.yaml @@ -37,7 +37,7 @@ social_security_cola_means_test_magi: 2026: 100001 output: - social_security_retirement_cola_withheld: [false, true, false] + social_security_retirement_cola_means_test_withheld: [false, true, false] # January 2027 statutory rate is 3.6%. The CPI-E addition appears only # in the two paid numerators; the withheld numerator is exactly one. social_security_retirement: diff --git a/policyengine_us/tests/unit/test_ss_cola_means_test.py b/policyengine_us/tests/unit/test_ss_cola_means_test.py index 6edd6e278f9..7a9890c8fb7 100644 --- a/policyengine_us/tests/unit/test_ss_cola_means_test.py +++ b/policyengine_us/tests/unit/test_ss_cola_means_test.py @@ -249,7 +249,9 @@ def test_each_index_scope_combination_and_first_eligible_boundary(index, scope): ) for component in COMPONENTS: np.testing.assert_array_equal( - simulation.calculate(f"social_security_{component}_cola_withheld", 2027), + simulation.calculate( + f"social_security_{component}_cola_means_test_withheld", 2027 + ), expected_withholding, ) np.testing.assert_allclose( @@ -275,7 +277,9 @@ def test_equality_is_paid_and_each_filing_status_has_an_explicit_group( system=means_test_system, ) np.testing.assert_array_equal( - simulation.calculate("social_security_retirement_cola_withheld", 2027), + simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2027 + ), [True, True, True, False, False, False], ) @@ -287,7 +291,9 @@ def test_surviving_spouse_mapping_is_a_parameter(): ), _reform(**{"filing_groups.SURVIVING_SPOUSE": "single"}), ) - assert simulation.calculate("social_security_retirement_cola_withheld", 2027)[0] + assert simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2027 + )[0] def test_static_cohort_proxy_is_reconstructed_within_each_score_year(means_test_system): @@ -328,7 +334,9 @@ def test_only_prior_tax_year_income_and_filing_status_make_the_decision( ), system=means_test_system, ) - assert simulation.calculate("social_security_retirement_cola_withheld", 2027)[0] + assert simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2027 + )[0] def test_magi_excludes_social_security_and_retains_irmaa_addbacks(): @@ -389,9 +397,9 @@ def test_supplied_income_history_overrides_the_extended_record_proxy( ) decisions.append( bool( - simulation.calculate("social_security_retirement_cola_withheld", 2027)[ - 0 - ] + simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2027 + )[0] ) ) assert decisions == [False, True] @@ -407,7 +415,11 @@ def test_static_cross_section_proxy_sensitivities(sensitivity, income, expected) _reform(income_sensitivity=sensitivity), ) assert ( - bool(simulation.calculate("social_security_retirement_cola_withheld", 2027)[0]) + bool( + simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2027 + )[0] + ) is expected ) @@ -415,12 +427,14 @@ def test_static_cross_section_proxy_sensitivities(sensitivity, income, expected) def test_explicit_withholding_input_remains_a_final_decision(means_test_system): simulation = _simulation( _situation( - {"worker": _record(social_security_retirement_cola_withheld={2027: True})}, + {"worker": _record(social_security_retirement_cola_override={2027: 1})}, {"worker": 0}, ), system=means_test_system, ) - assert simulation.calculate("social_security_retirement_cola_withheld", 2027)[0] + assert simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2027 + )[0] assert simulation.calculate("social_security_retirement", 2027)[0] < 12_000 @@ -430,7 +444,9 @@ def test_disabled_means_test_preserves_manual_input_carry_over(): {"worker": _record(social_security_retirement_cola_withheld={2026: True})} ) ) - assert simulation.calculate("social_security_retirement_cola_withheld", 2027)[0] + assert simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2027 + )[0] def test_zero_cola_and_downward_adjustment_floor(): @@ -492,8 +508,12 @@ def test_pre_receipt_index_is_separate_and_has_no_income_withholding( (1 + chained[year]) / (1 + baseline[year]) if use_chained_before_receipt else 1 for year in (2027, 2028, 2029) ) / (1 + baseline[2030]) - assert not simulation.calculate("social_security_retirement_cola_withheld", 2029)[0] - assert simulation.calculate("social_security_retirement_cola_withheld", 2030)[0] + assert not simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2029 + )[0] + assert simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2030 + )[0] assert simulation.calculate("social_security_retirement_joint_cola_factor", 2030)[ 0 ] == pytest.approx(expected) @@ -518,6 +538,111 @@ def _calculator_record(name): return record +@pytest.mark.parametrize("calculator", [False, True]) +def test_partial_paid_override_preserves_other_persons_automatic_decision( + means_test_system, calculator +): + # Core stores a complete array for each supplied variable/year. The -1 + # default must keep A automatic when only B supplies a paid decision. + records = { + name: _calculator_record(name) if calculator else _record() + for name in ("A", "B") + } + override = ( + "ss_calculator_cola_override" + if calculator + else "social_security_retirement_cola_override" + ) + records["B"][override] = {2027: 0} + simulation = _simulation( + _situation(records, {"A": 1_000_000, "B": 1_000_000}), + system=means_test_system, + ) + variable, period, expected = ( + ("ss_calculated_payment", "2027-01", [1000, 1036]) + if calculator + else ("social_security_retirement", 2027, [11583.01, 12000]) + ) + np.testing.assert_allclose( + simulation.calculate(variable, period), expected, rtol=0, atol=0.01 + ) + + +def test_calculator_override_precedence_and_exact_year_expiry(means_test_system): + # Exercise Core's separate input arrays: component decisions reach the + # calculator, and its own non-auto input takes precedence for that person. + records = {name: _calculator_record(name) for name in ("A", "B", "C")} + records["A"]["social_security_retirement_cola_override"] = {2027: 0} + records["B"].update( + social_security_retirement_cola_override={2027: 1}, + ss_calculator_cola_override={2027: 0}, + ) + records["C"].update( + social_security_retirement_cola_override={2027: 0}, + ss_calculator_cola_override={2027: 1}, + ) + simulation = _simulation( + _situation( + records, + {name: {2026: 1_000_000, 2027: 1_000_000} for name in records}, + ), + system=means_test_system, + ) + np.testing.assert_array_equal( + simulation.calculate("ss_calculated_payment", "2027-01"), + [1036, 1036, 1000], + ) + # All overrides expire after 2027. High prior-year income withholds the + # next COLA for all three; paid 2027 COLAs remain in A and B's check levels. + np.testing.assert_array_equal( + simulation.calculate("ss_calculated_payment", "2028-01"), + [1036, 1036, 1000], + ) + + +def test_disabled_means_test_calculator_ignores_component_withholding(): + # Literal checks captured from main 37dd5099ed7fad694419fb652eb3ba7bd14122ba. + # A5's calculator channel is exact-year ss_calculator_withhold_cola only. + record = _calculator_record("worker") + record["ss_calculator_supplied_pia"] = 1000.1 + record["social_security_retirement_cola_withheld"] = {2027: True} + simulation = _simulation( + _situation({"worker": record}), + reform={ + "gov.contrib.ssa.benefit_levers.cola.withholding_enabled": { + "2020-01-01": True + } + }, + ) + for year, expected in zip(range(2027, 2031), (1036, 1062, 1086, 1111)): + assert ( + simulation.calculate("ss_calculated_payment", f"{year}-01")[0] == expected + ) + + +def test_disabled_means_test_calculator_manual_decision_is_exact_year(): + records = {name: _calculator_record(name) for name in ("A", "B")} + for record in records.values(): + record["ss_calculator_supplied_pia"] = 1000.1 + records["B"]["ss_calculator_withhold_cola"] = {2027: True} + # New E3 inputs have no effect when E3 is off. + records["A"]["ss_calculator_cola_override"] = {2027: 1} + simulation = _simulation( + _situation(records), + reform={ + "gov.contrib.ssa.benefit_levers.cola.withholding_enabled": { + "2020-01-01": True + } + }, + ) + np.testing.assert_array_equal( + simulation.calculate("ss_calculated_payment", "2027-01"), [1036, 1000] + ) + np.testing.assert_array_equal( + simulation.calculate("ss_calculated_payment", "2028-01"), [1062, 1025] + ) + + @pytest.mark.parametrize("other_payment_factors", [False, True]) def test_calculated_source_applies_withholding_and_cpi_e_once(other_payment_factors): reform = _reform() @@ -587,11 +712,15 @@ def test_midyear_first_check_retains_the_pre_receipt_index(): # First receipt is July2028. January2028's COLA therefore uses the # pre-receipt index regardless of income; January2029 is income-tested. np.testing.assert_array_equal( - simulation.calculate("social_security_retirement_cola_withheld", 2028), + simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2028 + ), [False, False], ) np.testing.assert_array_equal( - simulation.calculate("social_security_retirement_cola_withheld", 2029), + simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2029 + ), [True, True], ) parameters = simulation.tax_benefit_system.parameters @@ -639,11 +768,15 @@ def test_named_withholding_uses_selected_calculator_cohort_and_receipt( simulation = _simulation(situation, system=means_test_system) baseline = _simulation(situation, system=disabled_means_test_system) np.testing.assert_array_equal( - simulation.calculate("social_security_retirement_cola_withheld", 2028), + simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2028 + ), [False, False], ) np.testing.assert_array_equal( - simulation.calculate("social_security_retirement_cola_withheld", 2029), + simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2029 + ), [False, True], ) np.testing.assert_array_equal( @@ -657,6 +790,16 @@ def test_named_withholding_uses_selected_calculator_cohort_and_receipt( def test_a9_preset_reproduces_captured_description(): + # Literal line boundaries pin the model's projection path independently + # of indexed_threshold's implementation. Equality is paid in both groups. + names = ( + "below", + "above", + "single_equal", + "single_above", + "joint_equal", + "joint_above", + ) simulation = _simulation( _situation( { @@ -665,9 +808,22 @@ def test_a9_preset_reproduces_captured_description(): age={2028: 70, 2029: 70}, social_security_retirement_first_cola_year=2029, ) - for name in ("below", "above") + for name in names + }, + { + "below": {2027: 0, 2028: 0}, + "above": {2027: 1_000_000, 2028: 1_000_000}, + "single_equal": {2028: 117_800}, + "single_above": {2028: 117_801}, + "joint_equal": {2028: 235_700}, + "joint_above": {2028: 235_701}, + }, + { + "single_equal": {2028: "SINGLE"}, + "single_above": {2028: "SINGLE"}, + "joint_equal": {2028: "JOINT"}, + "joint_above": {2028: "JOINT"}, }, - {"below": {2027: 0, 2028: 0}, "above": {2027: 1_000_000, 2028: 1_000_000}}, ), a9, ) @@ -682,15 +838,22 @@ def test_a9_preset_reproduces_captured_description(): assert means_test.thresholds.base_year == 2018 assert means_test.thresholds.single == 85_000 assert means_test.thresholds.joint == 170_000 + parameters = simulation.tax_benefit_system.parameters + index = parameters.gov.bls.cpi.cpi_u + assert index("2026-12-01") == 332.568 # Latest BLS value: June 2026. + assert index("2028-12-01") == 348.3 # CBO value authored February 2028. np.testing.assert_array_equal( - simulation.calculate("social_security_retirement_cola_withheld", 2028), - [False, False], + simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2028 + ), + [False] * len(names), ) np.testing.assert_array_equal( - simulation.calculate("social_security_retirement_cola_withheld", 2029), - [False, True], + simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2029 + ), + [False, True, False, True, False, True], ) - parameters = simulation.tax_benefit_system.parameters baseline = next( event["rate"] for event in cola_events(parameters, 2029) @@ -699,7 +862,7 @@ def test_a9_preset_reproduces_captured_description(): chained = chained_cola_rates(parameters, 2029)[2029] np.testing.assert_allclose( simulation.calculate("social_security_retirement_joint_cola_factor", 2029), - [(1 + chained) / (1 + baseline), 1 / (1 + baseline)], + [(1 + chained) / (1 + baseline), 1 / (1 + baseline)] * 3, ) @@ -714,7 +877,9 @@ def test_disabled_later_policy_retains_an_earlier_withheld_event(): for event in cola_events(simulation.tax_benefit_system.parameters, 2027) if event["payment_year"] == 2027 ) - assert not simulation.calculate("social_security_retirement_cola_withheld", 2028)[0] + assert not simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2028 + )[0] assert simulation.calculate("social_security_retirement_joint_cola_factor", 2028)[ 0 ] == pytest.approx(1 / (1 + baseline)) @@ -795,7 +960,9 @@ def test_property_withholding_is_monotone_in_income(means_test_system, lower, in ), system=means_test_system, ) - decisions = simulation.calculate("social_security_retirement_cola_withheld", 2027) + decisions = simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2027 + ) assert int(decisions[0]) <= int(decisions[1]) @@ -811,51 +978,77 @@ def test_property_a_withheld_year_never_raises_the_benefit(rate, delta): assert withheld <= paid -@settings(max_examples=60, deadline=None) -@given( - values=st.lists( - st.tuples(st.integers(0, 100), st.integers(-100, 100), st.booleans()), - min_size=1, - max_size=8, - ) -) -def test_property_joint_factor_equals_product_of_per_year_ratios(values): - events = [_event(2027 + i, c / 1000) for i, (c, _, _) in enumerate(values)] - deltas = {2027 + i: d / 10000 for i, (_, d, _) in enumerate(values)} - withheld = {2027 + i: w for i, (_, _, w) in enumerate(values)} - expected = prod( - (1 if w else 1 + max(0, c / 1000 + d / 10000)) / (1 + c / 1000) - for c, d, w in values - ) - actual = joint_cola_factor(events, (2027, 1), 2034, deltas, withheld) - assert actual == pytest.approx(expected) - - -@settings(max_examples=8, deadline=None) -@given( - reported=st.integers(0, 100_000), - prior_income=st.integers(0, 1_000_000), +@pytest.mark.parametrize( + "reported,prior_income,expected_components,expected_total", + [ + (0, 0, [0, 0, 0, 0], 0), + (12_000, 1_000_000, [12_000, 12_000, 12_000, 12_000], 48_000), + (100_000, 100_000, [100_000, 100_000, 100_000, 100_000], 400_000), + ], ) -def test_property_means_test_off_preserves_stacked_base_outputs( - disabled_means_test_system, reported, prior_income +def test_means_test_off_matches_main_stacked_base_outputs( + disabled_means_test_system, + reported, + prior_income, + expected_components, + expected_total, ): + # Literal outputs from main 37dd5099ed7fad694419fb652eb3ba7bd14122ba. + # Both public ingress modes must match main, even if both change together. record = _record() for component in COMPONENTS: record[f"social_security_{component}"] = {2027: reported} record.pop(f"social_security_{component}_reported") situation = _situation({"worker": record}, {"worker": prior_income}) - baseline = _simulation(situation) - disabled = _simulation(situation, system=disabled_means_test_system) - for variable in ( - *(f"social_security_{component}" for component in COMPONENTS), - "social_security", - *(f"social_security_{component}_joint_cola_factor" for component in COMPONENTS), - ): + for system in (None, disabled_means_test_system): + simulation = _simulation(situation, system=system) + for component, expected in zip(COMPONENTS, expected_components): + np.testing.assert_array_equal( + simulation.calculate(f"social_security_{component}", 2027), + [expected], + ) + np.testing.assert_array_equal( + simulation.calculate( + f"social_security_{component}_joint_cola_factor", 2027 + ), + [1], + ) np.testing.assert_array_equal( - disabled.calculate(variable, 2027), baseline.calculate(variable, 2027) + simulation.calculate("social_security", 2027), [expected_total] ) +def test_a5_only_joint_factor_ignores_calculator_population_dates(): + """Core must not validate irrelevant population dates on calculator rows.""" + calculated = _calculator_record("calculated") + calculated.update( + social_security_retirement_entitlement_year=2026, + social_security_retirement_entitlement_month=13, + ) + simulation = _simulation( + _situation( + { + "reported": _record( + social_security_retirement_cola_withheld={2027: True} + ), + "calculated": calculated, + } + ), + { + "gov.contrib.ssa.benefit_levers.cola.start_year": {"2020-01-01": 2027}, + "gov.contrib.ssa.benefit_levers.cola.withholding_enabled": { + "2020-01-01": True + }, + }, + ) + # The calculator's own dates are complete and valid. Its unused population + # month must not enter the A5 reported-route dependency graph when E3 is off. + np.testing.assert_allclose( + simulation.calculate("social_security_retirement_joint_cola_factor", 2027), + [0.9652509652509652, 1], + ) + + def _assert_cached_outputs_equal(reference, actual): for population in reference.populations.values(): for name, holder in population._holders.items(): @@ -873,6 +1066,19 @@ def _assert_cached_outputs_equal(reference, actual): def test_means_test_off_preserves_downstream_dependency_closure( disabled_means_test_system, calculated ): + # Main 37dd5099ed7fad694419fb652eb3ba7bd14122ba supplies these literals. + # Main cannot select the calculator as its public source, so the calculated + # case checks its direct calculator channel against main below and retains + # public downstream equality as supplemental disabled-policy coverage. + main_reported = { + "social_security": 24_000, + "taxable_social_security": 11_300, + "adjusted_gross_income": 41_300, + "income_tax": 2_715, + "net_investment_income_tax": 0, + "medicare_part_b_premium": 0, + "household_net_income": 49_442.6015625, + } record = _calculator_record("worker") if calculated else _record() if not calculated: # Differential ingress uses legacy public inputs, as M9 requires. @@ -899,13 +1105,23 @@ def test_means_test_off_preserves_downstream_dependency_closure( np.testing.assert_array_equal( disabled.calculate(name, 2027), baseline.calculate(name, 2027), err_msg=name ) + if not calculated: + for simulation in (baseline, disabled): + np.testing.assert_array_equal( + simulation.calculate(name, 2027), + [main_reported[name]], + err_msg=f"main: {name}", + ) if calculated: - for name, period in ( - ("ss_calculated_payment", "2027-01"), - ("ss_calculated_annual_entitlement", 2027), - ("ss_calculated_retirement", 2027), + for name, period, expected in ( + ("ss_calculated_payment", "2027-01", 1_036), + ("ss_calculated_annual_entitlement", 2027, 12_457), + ("ss_calculated_retirement", 2027, 12_432), ): - np.testing.assert_array_equal( - disabled.calculate(name, period), baseline.calculate(name, period) - ) + for simulation in (baseline, disabled): + np.testing.assert_array_equal( + simulation.calculate(name, period), + [expected], + err_msg=f"main: {name}", + ) _assert_cached_outputs_equal(baseline, disabled) diff --git a/policyengine_us/tools/social_security/engine.py b/policyengine_us/tools/social_security/engine.py index 0fc219e0c50..b49f055dab6 100644 --- a/policyengine_us/tools/social_security/engine.py +++ b/policyengine_us/tools/social_security/engine.py @@ -17,7 +17,7 @@ ) from .parameters import bend_points_at, formula_path_at from .pia import FormulaPath -from .means_test import component_withholding, indexed_cola_rate +from .means_test import cola_override, component_withholding, indexed_cola_rate def calculator_output(person, period, parameters, output): @@ -222,25 +222,31 @@ def money(name, when=period): withholding_enabled = ( policy.withholding_enabled or policy.means_test.enabled ) - if withholding_enabled and has_input_for_period( + if policy.means_test.enabled: + override = cola_override( + person, "ss_calculator_cola_override", str(payment_year) + )[index] + withheld = ( + bool(override) + if override != -1 + else bool( + component_withholding( + person, + payment_year, + parameters, + component, + eligibility_year=eligibility, + receipt_year=receipt, + receipt_month=receipt_month, + )[index] + ) + ) + elif withholding_enabled and has_input_for_period( person.simulation, "ss_calculator_withhold_cola", str(payment_year), ): withheld = bool(value("withhold_cola", str(payment_year))) - elif withholding_enabled: - withheld = bool( - component_withholding( - person, - payment_year, - parameters, - component, - eligibility_year=eligibility, - receipt_year=receipt, - receipt_month=receipt_month, - use_explicit=True, - )[index] - ) else: withheld = False index_rate = indexed_cola_rate( diff --git a/policyengine_us/tools/social_security/means_test.py b/policyengine_us/tools/social_security/means_test.py index 15f93f56c47..7fce5f30faa 100644 --- a/policyengine_us/tools/social_security/means_test.py +++ b/policyengine_us/tools/social_security/means_test.py @@ -297,6 +297,16 @@ def cohort_mask(settings, eligibility_year): return np.asarray(eligibility_year) >= settings.first_eligible_year +def cola_override(person, name, period): + """Read an exact-year tri-state input without carrying past overrides.""" + if not has_input_for_period(person.simulation, name, period): + return np.full(len(person.ids), -1, dtype=int) + override = person(name, period) + if np.any(~np.isin(override, (-1, 0, 1))): + raise ValueError("COLA override must be -1 (auto), 0 (paid), or 1 (withheld)") + return override + + def component_withholding( person, period, @@ -311,29 +321,20 @@ def component_withholding( The reported-factor caller supplies its score-year eligibility proxy for every historical event. It does not re-age a static row within one score. - Explicit per-event withholding inputs retain A5's override channel. + E3 overrides are separate tri-state inputs, so an omitted person remains + automatic even when another person's decision is supplied. The legacy + boolean input remains A5's channel when E3 is disabled. """ period = to_period(period) year = period.start.year p = policy_settings(parameters, year) - if use_explicit and has_input_for_period( - person.simulation, f"social_security_{component}_cola_withheld", period - ): + if not p.enabled: return person(f"social_security_{component}_cola_withheld", period) - if not policy_active(p, year): - # Before E3, this was a pure annual input. Preserve Core's carry-over - # of supplied manual decisions when the automatic policy is inactive. - # Formula caches must never become manual decision sources. - name = f"social_security_{component}_cola_withheld" - earlier = [ - when - for when in person.simulation._get_set_input_periods(name) - if when.start < period.start - and has_input_for_period(person.simulation, name, when) - ] - if earlier: - return person(name, max(earlier, key=lambda when: when.start)) - return np.zeros(len(person.ids), dtype=bool) + override = cola_override( + person, f"social_security_{component}_cola_override", period + ) + if not policy_active(p, year) or np.all(override != -1): + return override == 1 eligibility_year = ( component_eligibility_year(person, period, component) if eligibility_year is None @@ -349,13 +350,15 @@ def component_withholding( if receipt_month is None else np.asarray(receipt_month) ) - eligible = cohort_mask(p, eligibility_year) & ~before_receipt( - year, receipt_year, receipt_month + eligible = ( + cohort_mask(p, eligibility_year) + & ~before_receipt(year, receipt_year, receipt_month) + & (override == -1) ) if not np.any(eligible): - return eligible - return eligible & person.tax_unit( - "social_security_cola_means_test_withheld", period + return override == 1 + return (override == 1) | ( + eligible & person.tax_unit("social_security_cola_means_test_withheld", period) ) diff --git a/policyengine_us/variables/gov/ssa/social_security/calculator/inputs.py b/policyengine_us/variables/gov/ssa/social_security/calculator/inputs.py index 85c11d308f5..83fe288dad7 100644 --- a/policyengine_us/variables/gov/ssa/social_security/calculator/inputs.py +++ b/policyengine_us/variables/gov/ssa/social_security/calculator/inputs.py @@ -373,4 +373,15 @@ class ss_calculator_withhold_cola(Variable): definition_period = YEAR default_value = False label = "Social Security calculator withhold cola" - documentation = "Already-resolved COLA withholding decision for this check year, for example from an E3 policy." + documentation = ( + "Legacy A5 withholding decision for this exact check year when E3 is disabled." + ) + + +class ss_calculator_cola_override(Variable): + value_type = int + entity = Person + definition_period = YEAR + default_value = -1 + label = "Social Security calculator COLA means-test override" + documentation = "Exact check-year E3 decision: -1 automatic (default), 0 paid, 1 withheld. Takes precedence over the component override for this person." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_cola_withheld.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_cola_withheld.py index 5570c210b36..f564e5a11fd 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_cola_withheld.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_cola_withheld.py @@ -7,8 +7,25 @@ class social_security_dependents_cola_withheld(Variable): default_value = False entity = Person definition_period = YEAR - label = "Social Security dependents cola withheld" - documentation = "Automatic prior-year income decision when the COLA means test is enabled, with explicit per-year decision overrides." + label = "Social Security dependents COLA withheld under A5" + documentation = "Legacy manual decision when the means test is disabled; E3 uses the separate tri-state COLA override." + + +class social_security_dependents_cola_override(Variable): + value_type = int + default_value = -1 + entity = Person + definition_period = YEAR + label = "Social Security dependents COLA means-test override" + documentation = "Exact payment-year E3 decision: -1 automatic (default), 0 paid, 1 withheld. Omitted people remain automatic." + + +class social_security_dependents_cola_means_test_withheld(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security dependents resolved COLA means-test decision" + documentation = "Automatic prior-year income decision with a separate per-person tri-state override; returns the legacy manual decision with E3 disabled." def formula(person, period, parameters): return component_withholding(person, period, parameters, "dependents") diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_cola_withheld.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_cola_withheld.py index 65234ed823b..71aca5e160e 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_disability_cola_withheld.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_cola_withheld.py @@ -7,8 +7,25 @@ class social_security_disability_cola_withheld(Variable): default_value = False entity = Person definition_period = YEAR - label = "Social Security disability cola withheld" - documentation = "Automatic prior-year income decision when the COLA means test is enabled, with explicit per-year decision overrides." + label = "Social Security disability COLA withheld under A5" + documentation = "Legacy manual decision when the means test is disabled; E3 uses the separate tri-state COLA override." + + +class social_security_disability_cola_override(Variable): + value_type = int + default_value = -1 + entity = Person + definition_period = YEAR + label = "Social Security disability COLA means-test override" + documentation = "Exact payment-year E3 decision: -1 automatic (default), 0 paid, 1 withheld. Omitted people remain automatic." + + +class social_security_disability_cola_means_test_withheld(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security disability resolved COLA means-test decision" + documentation = "Automatic prior-year income decision with a separate per-person tri-state override; returns the legacy manual decision with E3 disabled." def formula(person, period, parameters): return component_withholding(person, period, parameters, "disability") diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_cola_withheld.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_cola_withheld.py index 1051300d454..acf310c6119 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_cola_withheld.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_cola_withheld.py @@ -7,8 +7,25 @@ class social_security_retirement_cola_withheld(Variable): default_value = False entity = Person definition_period = YEAR - label = "Social Security retirement cola withheld" - documentation = "Automatic prior-year income decision when the COLA means test is enabled, with explicit per-year decision overrides." + label = "Social Security retirement COLA withheld under A5" + documentation = "Legacy manual decision when the means test is disabled; E3 uses the separate tri-state COLA override." + + +class social_security_retirement_cola_override(Variable): + value_type = int + default_value = -1 + entity = Person + definition_period = YEAR + label = "Social Security retirement COLA means-test override" + documentation = "Exact payment-year E3 decision: -1 automatic (default), 0 paid, 1 withheld. Omitted people remain automatic." + + +class social_security_retirement_cola_means_test_withheld(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security retirement resolved COLA means-test decision" + documentation = "Automatic prior-year income decision with a separate per-person tri-state override; returns the legacy manual decision with E3 disabled." def formula(person, period, parameters): return component_withholding(person, period, parameters, "retirement") diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_cola_withheld.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_cola_withheld.py index 2f0fe18a6d4..33847a17087 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_cola_withheld.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_cola_withheld.py @@ -7,8 +7,25 @@ class social_security_survivors_cola_withheld(Variable): default_value = False entity = Person definition_period = YEAR - label = "Social Security survivors cola withheld" - documentation = "Automatic prior-year income decision when the COLA means test is enabled, with explicit per-year decision overrides." + label = "Social Security survivors COLA withheld under A5" + documentation = "Legacy manual decision when the means test is disabled; E3 uses the separate tri-state COLA override." + + +class social_security_survivors_cola_override(Variable): + value_type = int + default_value = -1 + entity = Person + definition_period = YEAR + label = "Social Security survivors COLA means-test override" + documentation = "Exact payment-year E3 decision: -1 automatic (default), 0 paid, 1 withheld. Omitted people remain automatic." + + +class social_security_survivors_cola_means_test_withheld(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security survivors resolved COLA means-test decision" + documentation = "Automatic prior-year income decision with a separate per-person tri-state override; returns the legacy manual decision with E3 disabled." def formula(person, period, parameters): return component_withholding(person, period, parameters, "survivors") From b6a70f9e6a4d1a3ce2731e307982fec45e4e303d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 30 Oct 2026 23:00:22 -0400 Subject: [PATCH 73/76] Isolate exact-year Social Security COLA override inputs --- changelog.d/ss-cola-means-test-e3.added.md | 2 +- docs/social-security-cola-means-test.md | 16 ++- policyengine_us/data/economic_assumptions.py | 19 +++- .../social_security_cola_means_test.yaml | 44 -------- .../social_security_retirement.yaml | 30 ++--- .../tests/unit/test_ss_cola_means_test.py | 104 +++++++++++++++--- policyengine_us/tools/general.py | 15 ++- 7 files changed, 144 insertions(+), 86 deletions(-) delete mode 100644 policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_cola_means_test.yaml diff --git a/changelog.d/ss-cola-means-test-e3.added.md b/changelog.d/ss-cola-means-test-e3.added.md index 0f034bc7f21..b281e2f0f2f 100644 --- a/changelog.d/ss-cola-means-test-e3.added.md +++ b/changelog.d/ss-cola-means-test-e3.added.md @@ -1 +1 @@ -Add an opt-in Social Security COLA means test using prior-year MAGI excluding Social Security, configurable indexed thresholds and filing groups, cohort or all-beneficiary scope, statutory or chained CPI-U COLAs, income-history sensitivities, and an A9 preset. Apply each withholding decision jointly with other COLA changes in reported benefits and the household calculator, with independent per-person overrides while preserving current-law defaults and the legacy A5 withholding input channels. Document the latest-authored-value CPI-U lookup used to index projected thresholds. +Add an opt-in Social Security COLA means test using prior-year MAGI excluding Social Security, configurable indexed thresholds and filing groups, cohort or all-beneficiary scope, statutory or chained CPI-U COLAs, income-history sensitivities, and an A9 preset. Apply each withholding decision jointly with other COLA changes in reported benefits and the household calculator, with independent per-person overrides while preserving current-law defaults and the legacy A5 withholding input channels. Exclude the new controls from Medicare IRMAA income-history inference and expire them in generated dataset years while preserving supplied decisions. Document the latest-authored-value CPI-U lookup used to index projected thresholds. diff --git a/docs/social-security-cola-means-test.md b/docs/social-security-cola-means-test.md index 0a28c2eb602..f8ef131268e 100644 --- a/docs/social-security-cola-means-test.md +++ b/docs/social-security-cola-means-test.md @@ -71,6 +71,18 @@ restrictions apply to automatic decisions; an explicit override supplies the final decision for that person and year. The tri-state overrides expire after their supplied year; later years return to automatic decisions. +These controls do not establish income history for Medicare IRMAA's lookback. +With reform parameters at their defaults, supplying them changes no benefit, +tax or Medicare premium. This exclusion is specific to the new controls; +legacy component withholding inputs retain their existing history behavior. +Dataset extension resets all five controls to `-1` in generated future rows, +while preserving decisions in every year the dataset actually supplied. + +Core stores these inputs as integers: floats truncate first (for example, +`0.5` becomes `0`). Integers outside `-1`, `0` and `1` raise `ValueError` when +an active exact-year decision reads them. A multi-year input such as +`year:2027:2` creates no exact annual decision in either year. + With the means test disabled, the reported route retains A5's boolean `social_security_X_cola_withheld` inputs and their legacy carry-forward behavior. The calculator reads only `ss_calculator_withhold_cola`, in its @@ -183,7 +195,9 @@ The default one-year lag uses the December 1, 2028 lookup for January 2029 checks. The A9 preset's $85,000/$170,000 base lines, December 2018 value `251.233`, and projected value `348.3` give $117,800/$235,700 after nearest-$100 half-up rounding. OACT's displayed $117,300/$234,600 lines use a different -CPI projection. The December 1 lookup, one-year payment lag and rounding are +CPI projection. This December lookup is a scoring assumption: a single filer +at $117,500 falls between PolicyEngine's and OACT's lines. The December 1 +lookup, one-year payment lag and rounding are documented modeling conventions: A9 specifies CPI-U indexing without prescribing these details. Pin both the authored index path and these settings for a score. diff --git a/policyengine_us/data/economic_assumptions.py b/policyengine_us/data/economic_assumptions.py index 630978bc736..21889d10953 100644 --- a/policyengine_us/data/economic_assumptions.py +++ b/policyengine_us/data/economic_assumptions.py @@ -20,6 +20,17 @@ "calibration.gov.cbo.income_by_source.adjusted_gross_income" ) +# These controls describe a decision for one supplied payment year, rather +# than a characteristic to project. Legacy COLA-withheld booleans retain +# their existing carry-forward behavior. +_EXACT_YEAR_COLA_OVERRIDES = ( + "social_security_retirement_cola_override", + "social_security_disability_cola_override", + "social_security_survivors_cola_override", + "social_security_dependents_cola_override", + "ss_calculator_cola_override", +) + MICRODATA_UPRATING_OVERRIDES = { "american_opportunity_credit": DEFAULT_MICRODATA_UPRATING, "cdcc_relevant_expenses": DEFAULT_MICRODATA_UPRATING, @@ -99,7 +110,10 @@ def extend_single_year_dataset( If ``end_year`` is not provided, it defaults to the latest year covered by the CPI-U parameter (gov.bls.cpi.cpi_u). - Variables without an uprating parameter are carried forward unchanged. + Variables without an uprating parameter are carried forward unchanged, + except the four component COLA overrides and ``ss_calculator_cola_override``. + Those exact-year decisions become automatic (``-1``) in generated rows; + the supplied base year and legacy COLA-withheld booleans are preserved. """ if system is None: from policyengine_us.system import system as _system @@ -122,6 +136,9 @@ def extend_single_year_dataset( for year in range(start_year + 1, end_year + 1): next_year = dataset.copy(deep=False) next_year.time_period = str(year) + for name in _EXACT_YEAR_COLA_OVERRIDES: + if name in next_year.person: + next_year.person[name] = -1 datasets.append(next_year) multi_year_dataset = USMultiYearDataset(datasets=datasets) diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_cola_means_test.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_cola_means_test.yaml deleted file mode 100644 index 445edaa873c..00000000000 --- a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_cola_means_test.yaml +++ /dev/null @@ -1,44 +0,0 @@ -- name: Prior-year income withholds only above the line in covered cohorts - period: 2027 - absolute_error_margin: 0.01 - input: - gov.contrib.ssa.benefit_levers.cola.means_test.enabled: true - gov.contrib.ssa.benefit_levers.cola.means_test.thresholds.single: 100000 - gov.contrib.ssa.benefit_levers.cola.means_test.thresholds.joint: 200000 - gov.contrib.ssa.benefit_levers.cola.start_year: 2027 - gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 - people: - equal: - age: 63 - social_security_retirement_worker_eligibility_year: 2026 - social_security_retirement_first_cola_year: 2027 - social_security_retirement_reported: 12000 - above: - age: 63 - social_security_retirement_worker_eligibility_year: 2026 - social_security_retirement_first_cola_year: 2027 - social_security_retirement_reported: 12000 - older_cohort: - age: 64 - social_security_retirement_worker_eligibility_year: 2025 - social_security_retirement_first_cola_year: 2027 - social_security_retirement_reported: 12000 - tax_units: - equal: - members: [equal] - social_security_cola_means_test_magi: - 2026: 100000 - above: - members: [above] - social_security_cola_means_test_magi: - 2026: 100001 - older_cohort: - members: [older_cohort] - social_security_cola_means_test_magi: - 2026: 100001 - output: - social_security_retirement_cola_means_test_withheld: [false, true, false] - # January 2027 statutory rate is 3.6%. The CPI-E addition appears only - # in the two paid numerators; the withheld numerator is exactly one. - social_security_retirement: - [12017.374517374517, 11583.011583011583, 12017.374517374517] diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement.yaml index 38836902520..bb00b433829 100644 --- a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement.yaml +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement.yaml @@ -38,21 +38,21 @@ output: social_security_retirement: 12000 -- name: OASI-only payable share preserves disability benefits +- name: OASI-only payable share preserves disability and dependent DI benefits period: 2027 input: gov.contrib.ssa.benefit_levers.payable.share: 0.8 gov.contrib.ssa.benefit_levers.payable.oasi_only: true - social_security_dependents_fund: oasi - social_security_retirement_reported: 12000 - social_security_disability_reported: 6000 - social_security_survivors_reported: 3000 - social_security_dependents_reported: 1000 + social_security_dependents_fund: [oasi, di] + social_security_retirement_reported: [12000, 0] + social_security_disability_reported: [6000, 0] + social_security_survivors_reported: [3000, 0] + social_security_dependents_reported: [1000, 6000] output: - social_security_retirement: 9600 - social_security_disability: 6000 - social_security_survivors: 2400 - social_security_dependents: 800 + social_security_retirement: [9600, 0] + social_security_disability: [6000, 0] + social_security_survivors: [2400, 0] + social_security_dependents: [800, 6000] - name: A full across-the-board cut produces zero period: 2027 @@ -89,16 +89,6 @@ output: social_security_retirement: 10575.289575289577 -- name: OASI-only payable scenario preserves a dependent DI benefit - period: 2027 - input: - gov.contrib.ssa.benefit_levers.payable.share: 0.8 - gov.contrib.ssa.benefit_levers.payable.oasi_only: true - social_security_dependents_fund: di - social_security_dependents_reported: 6000 - output: - social_security_dependents: 6000 - - name: Complete calculated base bypasses every population benefit factor period: 2027 input: diff --git a/policyengine_us/tests/unit/test_ss_cola_means_test.py b/policyengine_us/tests/unit/test_ss_cola_means_test.py index 7a9890c8fb7..a8007f7cd0a 100644 --- a/policyengine_us/tests/unit/test_ss_cola_means_test.py +++ b/policyengine_us/tests/unit/test_ss_cola_means_test.py @@ -13,6 +13,7 @@ import pytest from hypothesis import given, settings, strategies as st from policyengine_core.parameters import ParameterNode +from policyengine_core.simulations import SimulationBuilder from policyengine_us import Simulation from policyengine_us.reforms.ssa.cola_means_test import a9 @@ -121,6 +122,17 @@ def _simulation(situation, reform=None, system=None): ) +def _override_situation(records, incomes=None): + """Exercise person decisions within one tax unit and across tax units.""" + situation = _situation(records, incomes) + situation["people"]["B"].update( + is_tax_unit_head=False, is_tax_unit_spouse=True + ) + situation["tax_units"]["A"]["members"] = ["A", "B"] + situation["tax_units"].pop("B") + return situation + + @pytest.fixture(scope="module") def means_test_system(): return _simulation(_situation(), _reform()).tax_benefit_system @@ -425,17 +437,24 @@ def test_static_cross_section_proxy_sensitivities(sensitivity, income, expected) def test_explicit_withholding_input_remains_a_final_decision(means_test_system): + records = {name: _record() for name in ("A", "B", "C")} + records["B"]["social_security_retirement_cola_override"] = {2027: 1} simulation = _simulation( - _situation( - {"worker": _record(social_security_retirement_cola_override={2027: 1})}, - {"worker": 0}, - ), + _override_situation(records, {name: 0 for name in records}), system=means_test_system, ) - assert simulation.calculate( - "social_security_retirement_cola_means_test_withheld", 2027 - )[0] - assert simulation.calculate("social_security_retirement", 2027)[0] < 12_000 + np.testing.assert_array_equal( + simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2027 + ), + [False, True, False], + ) + np.testing.assert_allclose( + simulation.calculate("social_security_retirement", 2027), + [12000, 11583.01, 12000], + rtol=0, + atol=0.01, + ) def test_disabled_means_test_preserves_manual_input_carry_over(): @@ -543,10 +562,11 @@ def test_partial_paid_override_preserves_other_persons_automatic_decision( means_test_system, calculator ): # Core stores a complete array for each supplied variable/year. The -1 - # default must keep A automatic when only B supplies a paid decision. + # default must keep A automatic when B in the same tax unit supplies a + # paid decision. C in another tax unit also remains automatic. records = { name: _calculator_record(name) if calculator else _record() - for name in ("A", "B") + for name in ("A", "B", "C") } override = ( "ss_calculator_cola_override" @@ -555,13 +575,13 @@ def test_partial_paid_override_preserves_other_persons_automatic_decision( ) records["B"][override] = {2027: 0} simulation = _simulation( - _situation(records, {"A": 1_000_000, "B": 1_000_000}), + _override_situation(records, {name: 1_000_000 for name in records}), system=means_test_system, ) variable, period, expected = ( - ("ss_calculated_payment", "2027-01", [1000, 1036]) + ("ss_calculated_payment", "2027-01", [1000, 1036, 1000]) if calculator - else ("social_security_retirement", 2027, [11583.01, 12000]) + else ("social_security_retirement", 2027, [11583.01, 12000, 11583.01]) ) np.testing.assert_allclose( simulation.calculate(variable, period), expected, rtol=0, atol=0.01 @@ -582,7 +602,7 @@ def test_calculator_override_precedence_and_exact_year_expiry(means_test_system) ss_calculator_cola_override={2027: 1}, ) simulation = _simulation( - _situation( + _override_situation( records, {name: {2026: 1_000_000, 2027: 1_000_000} for name in records}, ), @@ -621,14 +641,14 @@ def test_disabled_means_test_calculator_ignores_component_withholding(): def test_disabled_means_test_calculator_manual_decision_is_exact_year(): - records = {name: _calculator_record(name) for name in ("A", "B")} + records = {name: _calculator_record(name) for name in ("A", "B", "C")} for record in records.values(): record["ss_calculator_supplied_pia"] = 1000.1 records["B"]["ss_calculator_withhold_cola"] = {2027: True} # New E3 inputs have no effect when E3 is off. records["A"]["ss_calculator_cola_override"] = {2027: 1} simulation = _simulation( - _situation(records), + _override_situation(records), reform={ "gov.contrib.ssa.benefit_levers.cola.withholding_enabled": { "2020-01-01": True @@ -636,10 +656,10 @@ def test_disabled_means_test_calculator_manual_decision_is_exact_year(): }, ) np.testing.assert_array_equal( - simulation.calculate("ss_calculated_payment", "2027-01"), [1036, 1000] + simulation.calculate("ss_calculated_payment", "2027-01"), [1036, 1000, 1036] ) np.testing.assert_array_equal( - simulation.calculate("ss_calculated_payment", "2028-01"), [1062, 1025] + simulation.calculate("ss_calculated_payment", "2028-01"), [1062, 1025, 1062] ) @@ -689,6 +709,54 @@ def test_calculated_source_applies_withholding_and_cpi_e_once(other_payment_fact simulation.calculate("social_security_retirement", 2027), [10_752, 10_368] if other_payment_factors else [12_444, 12_000], ) + if not other_payment_factors: + _assert_raw_builder_means_test(simulation.tax_benefit_system) + + +def _assert_raw_builder_means_test(system): + # Preserve the former YAML case's raw Core ingress while reusing the + # existing reform system, avoiding a separate policy CLI process. + builder = SimulationBuilder() + builder.set_default_period("2027") + simulation = builder.build_from_dict( + system, + { + "people": { + name: { + "age": 64 if name == "older_cohort" else 63, + "social_security_retirement_worker_eligibility_year": ( + 2025 if name == "older_cohort" else 2026 + ), + "social_security_retirement_first_cola_year": 2027, + "social_security_retirement_reported": 12000, + } + for name in ("equal", "above", "older_cohort") + }, + "tax_units": { + name: { + "members": [name], + "social_security_cola_means_test_magi": { + 2026: 100000 if name == "equal" else 100001 + }, + } + for name in ("equal", "above", "older_cohort") + }, + }, + ) + np.testing.assert_array_equal( + simulation.calculate( + "social_security_retirement_cola_means_test_withheld", 2027 + ), + [False, True, False], + ) + # January 2027 statutory rate is 3.6%. The CPI-E addition appears only + # in the two paid numerators; the withheld numerator is exactly one. + np.testing.assert_allclose( + simulation.calculate("social_security_retirement", 2027), + [12017.374517374517, 11583.011583011583, 12017.374517374517], + rtol=0, + atol=0.01, + ) def test_midyear_first_check_retains_the_pre_receipt_index(): diff --git a/policyengine_us/tools/general.py b/policyengine_us/tools/general.py index a5022e0f76b..7afc7860fd7 100644 --- a/policyengine_us/tools/general.py +++ b/policyengine_us/tools/general.py @@ -56,7 +56,13 @@ def earliest_input_year(simulation): neutralized, are skipped: they are never read. Social Security calculator records and output overrides are also skipped: a historical earnings-test amount, COLA decision or calculator result is not a tax-year income return - for the IRMAA lookback. As with + for the IRMAA lookback. The four component controls + ``social_security_retirement_cola_override``, + ``social_security_disability_cola_override``, + ``social_security_survivors_cola_override`` and + ``social_security_dependents_cola_override`` are excluded by exact name + for the same reason. Legacy ``social_security_*_cola_withheld`` inputs + still count, preserving their existing income-history behavior. As with ``has_input_for_period``, core versions before 3.32.27 can keep an input key after the input is deleted, and a value recalculated for that period then still counts. Formulas that read this cache their @@ -70,6 +76,13 @@ def earliest_input_year(simulation): branch_name in visible_branches and period.unit != ETERNITY and not variable.startswith(("ss_calculator_", "ss_calculated_")) + and variable + not in ( + "social_security_retirement_cola_override", + "social_security_disability_cola_override", + "social_security_survivors_cola_override", + "social_security_dependents_cola_override", + ) and variable in variables and not variables[variable].is_neutralized ): From 53c03b7c00aa69de48d317f681abb1709b7992cb Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 30 Oct 2026 23:01:51 -0400 Subject: [PATCH 74/76] Test COLA controls across default routes and dataset projection --- .../unit/test_ss_cola_dataset_extension.py | 161 ++++++++++++++++++ .../unit/test_ss_cola_default_invariance.py | 155 +++++++++++++++++ .../tests/unit/test_ss_cola_means_test.py | 4 +- 3 files changed, 317 insertions(+), 3 deletions(-) create mode 100644 policyengine_us/tests/unit/test_ss_cola_dataset_extension.py create mode 100644 policyengine_us/tests/unit/test_ss_cola_default_invariance.py diff --git a/policyengine_us/tests/unit/test_ss_cola_dataset_extension.py b/policyengine_us/tests/unit/test_ss_cola_dataset_extension.py new file mode 100644 index 00000000000..6ca7b03b03a --- /dev/null +++ b/policyengine_us/tests/unit/test_ss_cola_dataset_extension.py @@ -0,0 +1,161 @@ +"""Exact-year COLA overrides expire only in generated dataset years.""" + +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest + +from policyengine_us import Simulation +from policyengine_us.data import USMultiYearDataset, USSingleYearDataset +from policyengine_us.data.economic_assumptions import ( + _apply_uprating, + extend_single_year_dataset, +) +from policyengine_us.tests.unit.test_ss_cola_means_test import ( + COMPONENTS, + _calculator_record, + _record, + _reform, + _simulation, + _situation, +) + + +OVERRIDES = tuple( + f"social_security_{component}_cola_override" for component in COMPONENTS +) + ("ss_calculator_cola_override",) +LEGACY_WITHHELD = tuple( + f"social_security_{component}_cola_withheld" for component in COMPONENTS +) + ("ss_calculator_withhold_cola",) + + +def _controls_dataset(year=2027): + return USSingleYearDataset( + person=pd.DataFrame( + { + "person_id": [1, 2, 3], + **{name: [-1, 0, 1] for name in OVERRIDES}, + **{name: [False, True, False] for name in LEGACY_WITHHELD}, + } + ), + household=pd.DataFrame({"household_id": [1]}), + tax_unit=pd.DataFrame({"tax_unit_id": [1]}), + time_period=year, + ) + + +def _assert_generated_controls(dataset): + for name in OVERRIDES: + np.testing.assert_array_equal(dataset[2027].person[name], [-1, 0, 1]) + for year in (2028, 2029): + np.testing.assert_array_equal(dataset[year].person[name], [-1, -1, -1]) + for year in dataset.years: + for name in LEGACY_WITHHELD: + np.testing.assert_array_equal( + dataset[year].person[name], [False, True, False] + ) + + +def test_generated_years_reset_all_five_overrides_and_preserve_legacy_booleans(): + source = _controls_dataset() + original = source.person.copy() + extended = extend_single_year_dataset( + source, end_year=2029, system=SimpleNamespace(variables={}) + ) + _assert_generated_controls(extended) + pd.testing.assert_frame_equal(source.person, original) + + +def test_supplied_multiyear_decisions_survive_uprating_and_serialization(tmp_path): + first = _controls_dataset() + second = _controls_dataset(2028) + for name in OVERRIDES: + second.person[name] = [1, -1, 0] + supplied = USMultiYearDataset(datasets=[first, second]) + uprated = _apply_uprating(supplied, system=SimpleNamespace(variables={})) + path = str(tmp_path / "supplied_years.h5") + uprated.save(path) + restored = USMultiYearDataset(file_path=path) + for year in supplied.years: + pd.testing.assert_frame_equal(restored[year].person, supplied[year].person) + + +@pytest.mark.parametrize("from_file", [False, True]) +def test_constructor_generation_routes_expire_overrides_without_running_microsim( + from_file, tmp_path, monkeypatch +): + # Exercise both constructor routes through the real extender. Stub Core's + # constructor so this routing test neither loads nor runs a microsimulation. + from policyengine_us import system as system_module + from policyengine_us.tests.microsimulation.data.fixtures.test_extend_single_year_dataset import ( + make_mock_super_init, + ) + + source = _controls_dataset() + dataset = source + if from_file: + dataset = str(tmp_path / "single_year.h5") + source.save(dataset) + captured = [None] + monkeypatch.setattr( + system_module.CoreMicrosimulation, + "__init__", + make_mock_super_init(system_module, captured=captured), + ) + system_module.Microsimulation(dataset=dataset, dataset_end_year=2029) + _assert_generated_controls(captured[0]) + + +@pytest.fixture(scope="module") +def means_test_system(): + return _simulation(_situation(), _reform()).tax_benefit_system + + +@pytest.mark.parametrize("calculator", [False, True]) +def test_generated_year_does_not_repeat_a_paid_cola_override( + means_test_system, calculator +): + record = _calculator_record("worker") if calculator else _record() + row = { + name: value[2027] if isinstance(value, dict) else value + for name, value in record.items() + } + entities = ("household", "tax_unit", "spm_unit", "family", "marital_unit") + row.update(person_id=1, **{f"person_{entity}_id": 1 for entity in entities}) + override = ( + "ss_calculator_cola_override" + if calculator + else "social_security_retirement_cola_override" + ) + row[override] = 0 + source = USSingleYearDataset( + person=pd.DataFrame([row]), + household=pd.DataFrame({"household_id": [1], "household_weight": [1.0]}), + tax_unit=pd.DataFrame( + {"tax_unit_id": [1], "social_security_cola_means_test_magi": [1_000_000]} + ), + spm_unit=pd.DataFrame({"spm_unit_id": [1]}), + family=pd.DataFrame({"family_id": [1]}), + marital_unit=pd.DataFrame({"marital_unit_id": [1]}), + time_period=2027, + ) + extended = extend_single_year_dataset( + source, end_year=2028, system=means_test_system + ) + if calculator: + # PIA is a fixed eligibility-year source. Keep it fixed here so generic + # monetary uprating cannot change the separate COLA-expiry comparison. + for year in extended.years: + extended[year].person["ss_calculator_supplied_pia"] = 1000 + simulation = Simulation(dataset=extended, tax_benefit_system=means_test_system) + for year in (2026, 2027): + simulation.set_input("social_security_cola_means_test_magi", year, [1_000_000]) + variable, period, expected = ( + ("ss_calculated_payment", "2028-01", 1036) + if calculator + else ("social_security_retirement", 2028, 12352.54) + ) + np.testing.assert_allclose( + simulation.calculate(variable, period), [expected], rtol=0, atol=0.01 + ) diff --git a/policyengine_us/tests/unit/test_ss_cola_default_invariance.py b/policyengine_us/tests/unit/test_ss_cola_default_invariance.py new file mode 100644 index 00000000000..8cd77c86ff9 --- /dev/null +++ b/policyengine_us/tests/unit/test_ss_cola_default_invariance.py @@ -0,0 +1,155 @@ +"""New reform controls must be inert under current law, including input history. + +Python covers Core's input provenance, constructor migration, and generated +years/values. All cases reuse the default country system without a reform. +""" + +from copy import deepcopy + +import numpy as np +import pytest +from hypothesis import example, given, settings, strategies as st + +from policyengine_us.tests.unit.test_ss_cola_means_test import ( + COMPONENTS, + _calculator_record, + _record, + _simulation, + _situation, +) +from policyengine_us.tools.general import earliest_input_year + + +OVERRIDES = tuple(f"social_security_{c}_cola_override" for c in COMPONENTS) + ( + "ss_calculator_cola_override", +) + + +@pytest.fixture(scope="module") +def irmaa_control(): + situation = _situation( + { + "worker": { + "age": {2027: 70}, + "employment_income": {2027: 500_000}, + "is_tax_unit_head": True, + } + } + ) + # The review's minimal record has no auxiliary tax-unit input history. + situation["tax_units"]["worker"].pop("tax_unit_itemizes") + simulation = _simulation(situation) + values = { + name: simulation.calculate(name, 2027).copy() + for name in ("medicare_irmaa_magi_two_years_prior", "medicare_part_b_premium") + } + np.testing.assert_allclose( + values["medicare_irmaa_magi_two_years_prior"], [463100.31], rtol=0, atol=0.01 + ) + np.testing.assert_allclose( + values["medicare_part_b_premium"], [7790.40], rtol=0, atol=0.01 + ) + return situation, values + + +@pytest.mark.parametrize("name", OVERRIDES) +@pytest.mark.parametrize("value", [-1, 1]) +def test_override_does_not_supply_irmaa_income_history(irmaa_control, name, value): + situation, expected = irmaa_control + situation = deepcopy(situation) + situation["people"]["worker"][name] = {2020: value} + simulation = _simulation(situation) + assert earliest_input_year(simulation) == 2027 + for output, control in expected.items(): + np.testing.assert_array_equal(simulation.calculate(output, 2027), control) + + +def test_legacy_withholding_still_supplies_input_history(irmaa_control): + situation, _ = irmaa_control + situation = deepcopy(situation) + for component in COMPONENTS: + situation["people"]["worker"][f"social_security_{component}_cola_withheld"] = { + 2020: False + } + assert earliest_input_year(_simulation(situation)) == 2020 + + +@pytest.fixture(scope="module") +def default_routes(): + reported = _record() + public = {name.removesuffix("_reported"): value for name, value in reported.items()} + records = { + "public": public, + "reported": reported, + "calculator": _calculator_record("calculator"), + } + for record in records.values(): + record["employment_income"] = {2027: 500_000} + situation = _situation(records) + simulation = _simulation(situation) + outputs = ( + tuple(f"social_security_{c}" for c in COMPONENTS) + + tuple(f"social_security_{c}_cola_means_test_withheld" for c in COMPONENTS) + + tuple(f"social_security_{c}_joint_cola_factor" for c in COMPONENTS) + + ( + "social_security", + "taxable_social_security", + "adjusted_gross_income", + "income_tax", + "ssi", + "household_net_income", + "medicare_irmaa_magi_two_years_prior", + "medicare_part_b_premium", + "ss_calculated_annual_entitlement", + "ss_calculated_annual_payment", + ) + ) + expected = { + (name, 2027): simulation.calculate(name, 2027).copy() for name in outputs + } + expected.update( + { + ("ss_calculated_payment", month): simulation.calculate( + "ss_calculated_payment", month + ).copy() + for month in ("2027-01", "2028-01") + } + ) + return situation, expected, earliest_input_year(simulation) + + +@pytest.mark.parametrize("ingress", ["situation", "set_input"]) +@settings(max_examples=6, deadline=None) +@example(name=OVERRIDES[0], year=2027, value=2) +@example(name=OVERRIDES[1], year=2027, value=-2) +@example(name=OVERRIDES[2], year=2027, value=0.5) +@example(name=OVERRIDES[3], year=2027, value=1) +@example(name=OVERRIDES[4], year=2027, value=2) +@given( + name=st.sampled_from(OVERRIDES), + year=st.integers(min_value=1900, max_value=2100), + value=st.one_of( + st.integers(-10_000, 10_000), + st.floats(-10_000, 10_000, allow_nan=False, allow_infinity=False), + ), +) +def test_default_policy_ignores_override_years_and_values( + default_routes, ingress, name, year, value +): + """Any stored decision is inert on public, reported and calculator routes. + + Vary one control at a time so another channel cannot mask its effect. + Include invalid active-policy decisions: disabled policy must never read + or validate them. Raw input variables themselves are controls, not outputs. + """ + situation, expected, expected_earliest_year = default_routes + situation = deepcopy(situation) + if ingress == "situation": + for record in situation["people"].values(): + record[name] = {year: value} + simulation = _simulation(situation) + if ingress == "set_input": + simulation.set_input(name, year, np.full(len(situation["people"]), value)) + assert earliest_input_year(simulation) == expected_earliest_year + for (output, period), control in expected.items(): + np.testing.assert_array_equal(simulation.calculate(output, period), control) diff --git a/policyengine_us/tests/unit/test_ss_cola_means_test.py b/policyengine_us/tests/unit/test_ss_cola_means_test.py index a8007f7cd0a..e43f3dafbbf 100644 --- a/policyengine_us/tests/unit/test_ss_cola_means_test.py +++ b/policyengine_us/tests/unit/test_ss_cola_means_test.py @@ -125,9 +125,7 @@ def _simulation(situation, reform=None, system=None): def _override_situation(records, incomes=None): """Exercise person decisions within one tax unit and across tax units.""" situation = _situation(records, incomes) - situation["people"]["B"].update( - is_tax_unit_head=False, is_tax_unit_spouse=True - ) + situation["people"]["B"].update(is_tax_unit_head=False, is_tax_unit_spouse=True) situation["tax_units"]["A"]["members"] = ["A", "B"] situation["tax_units"].pop("B") return situation From b33afe8dba0ac6a948500d1ba096c4e62528e61d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 30 Oct 2026 23:13:23 -0400 Subject: [PATCH 75/76] Date the Medicare regression fixture explicitly --- policyengine_us/tests/unit/test_ss_cola_default_invariance.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/policyengine_us/tests/unit/test_ss_cola_default_invariance.py b/policyengine_us/tests/unit/test_ss_cola_default_invariance.py index 8cd77c86ff9..8bfb483a146 100644 --- a/policyengine_us/tests/unit/test_ss_cola_default_invariance.py +++ b/policyengine_us/tests/unit/test_ss_cola_default_invariance.py @@ -32,13 +32,14 @@ def irmaa_control(): "worker": { "age": {2027: 70}, "employment_income": {2027: 500_000}, - "is_tax_unit_head": True, + "is_tax_unit_head": {2027: True}, } } ) # The review's minimal record has no auxiliary tax-unit input history. situation["tax_units"]["worker"].pop("tax_unit_itemizes") simulation = _simulation(situation) + assert earliest_input_year(simulation) == 2027 values = { name: simulation.calculate(name, 2027).copy() for name in ("medicare_irmaa_magi_two_years_prior", "medicare_part_b_premium") From 0a6dedd02abec864969f0444f612d216989733af Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 1 Nov 2026 01:38:02 -0400 Subject: [PATCH 76/76] Guard COLA override inventory for reform scoring --- .../unit/test_ss_cola_default_invariance.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/policyengine_us/tests/unit/test_ss_cola_default_invariance.py b/policyengine_us/tests/unit/test_ss_cola_default_invariance.py index 8bfb483a146..54284aef4ac 100644 --- a/policyengine_us/tests/unit/test_ss_cola_default_invariance.py +++ b/policyengine_us/tests/unit/test_ss_cola_default_invariance.py @@ -5,11 +5,15 @@ """ from copy import deepcopy +from types import SimpleNamespace import numpy as np import pytest from hypothesis import example, given, settings, strategies as st +from policyengine_core.periods import period as to_period +from policyengine_us import Simulation +from policyengine_us.data.economic_assumptions import _EXACT_YEAR_COLA_OVERRIDES from policyengine_us.tests.unit.test_ss_cola_means_test import ( COMPONENTS, _calculator_record, @@ -25,6 +29,41 @@ ) +def test_cola_override_inventory_covers_inference_reset_and_invariance(): + system = Simulation.default_tax_benefit_system_instance + registered = {name for name in system.variables if name.endswith("_cola_override")} + coverage_message = ( + "New COLA override controls need inference, reset and invariance coverage: " + "update earliest_input_year, _EXACT_YEAR_COLA_OVERRIDES and OVERRIDES." + ) + assert registered == set(OVERRIDES), coverage_message + assert registered == set(_EXACT_YEAR_COLA_OVERRIDES), coverage_message + + # Exercise the exclusions without building another country system or simulation. + historical_period = to_period(2020) + holder = SimpleNamespace( + get_known_branch_periods=lambda: {("default", historical_period)} + ) + probe = SimpleNamespace( + tax_benefit_system=system, + _get_visible_branch_names=lambda: ["default"], + _user_input_keys=[(name, "default", historical_period) for name in registered], + get_holder=lambda name: holder, + ) + assert earliest_input_year(probe) is None, coverage_message + + # Component exclusions stay exact; calculator records use the prefix rule. + variables = system.variables.copy() + probe.tax_benefit_system = SimpleNamespace(variables=variables) + for name, expected in ( + ("social_security_auxiliary_cola_override", 2020), + ("ss_calculator_inventory_guard", None), + ): + variables[name] = system.variables[OVERRIDES[0]] + probe._user_input_keys = [(name, "default", historical_period)] + assert earliest_input_year(probe) == expected, name + + @pytest.fixture(scope="module") def irmaa_control(): situation = _situation(