From f68832e694c5bdc029005fbaead90ad3e658135f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 01:33:48 -0400 Subject: [PATCH 01/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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 b5656192351f763b7350fd12c7acb12efb50b145 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 27 Oct 2026 05:23:44 -0700 Subject: [PATCH 65/80] Add baseline-only reported benefit PIA ratio inference --- .../tests/unit/test_ss_pia_ratio.py | 298 ++++++++++++++++++ .../tools/social_security/pia_ratio.py | 208 ++++++++++++ 2 files changed, 506 insertions(+) create mode 100644 policyengine_us/tests/unit/test_ss_pia_ratio.py create mode 100644 policyengine_us/tools/social_security/pia_ratio.py diff --git a/policyengine_us/tests/unit/test_ss_pia_ratio.py b/policyengine_us/tests/unit/test_ss_pia_ratio.py new file mode 100644 index 00000000000..d7eb847b59b --- /dev/null +++ b/policyengine_us/tests/unit/test_ss_pia_ratio.py @@ -0,0 +1,298 @@ +"""Pure population ratio API: no simulation or parameter-tree construction.""" + +from decimal import Decimal +from fractions import Fraction + +from hypothesis import given, strategies as st +import pytest + +from policyengine_us.tools.social_security.pia import ( + FormulaPath, + compute_pia, + inverse_pia, +) +from policyengine_us.tools.social_security.pia_ratio import compute_pia_ratio + + +D = Decimal +BENDS = (960, 5785) +REVIEW_PATH = FormulaPath( + factor_values=(".95", ".32", ".15"), + bend_point_multipliers=("1.15", 1), +) + + +def test_c2_exact_interval_and_gain_before_and_after_dime_truncation(): + result = compute_pia_ratio( + "18117.12", + 2020, + REVIEW_PATH, + bend_points=BENDS, + claiming_adjustment=".70", + exact_eligibility_pia=True, + ) + assert result.aime == result.aime_lower == D(5000) + assert result.aime_upper == D("5000.3125") + assert not result.measurement_uncertainty + assert D("18117.12") * result.factor == D("19282.20") + before_rounding = compute_pia( + 5000, 2020, REVIEW_PATH, bend_points=BENDS, round_result=False + ) + baseline = compute_pia(5000, 2020, bend_points=BENDS, round_result=False) + assert 12 * D(".70") * (before_rounding - baseline) == D("1165.248") + + +def test_c2_fra_assumption_changes_inferred_aime_and_gain(): + result = compute_pia_ratio( + "18117.12", 2020, REVIEW_PATH, bend_points=BENDS, claiming_adjustment=1 + ) + assert result.aime == D(2978) + assert result.measurement_uncertainty + baseline = compute_pia(2978, 2020, bend_points=BENDS, round_result=False) + reform = compute_pia(2978, 2020, REVIEW_PATH, bend_points=BENDS, round_result=False) + assert 12 * (reform - baseline) == D("1664.64") + assert (D("1664.64") / D("1165.248") - 1) == pytest.approx(D(3) / 7) + + +def test_s770_eighteen_percent_bend_is_distinct_from_review_fixture(): + baseline = compute_pia(5000, 2020, bend_points=BENDS, round_result=False) + reform = compute_pia( + 5000, + 2020, + FormulaPath(s770=True), + 2026, + "retroactive", + bend_points=BENDS, + round_result=False, + ) + assert reform - baseline == D(".1634") * BENDS[0] + result = compute_pia_ratio( + "18117.12", + 2020, + FormulaPath(s770=True), + 2026, + "retroactive", + bend_points=BENDS, + claiming_adjustment=".7", + exact_eligibility_pia=True, + ) + assert D("18117.12") * result.factor == D("19434.24") + + +@pytest.mark.parametrize( + "pia", [".1", "863.9", "864", "864.1", "2407.9", "2408", "2408.1", "3000"] +) +def test_every_inverse_segment_and_boundary_uses_shared_lower_endpoint(pia): + result = compute_pia_ratio( + 12 * D(pia), 2020, bend_points=BENDS, exact_eligibility_pia=True + ) + interval = inverse_pia(pia, 2020, bend_points=BENDS) + assert result.aime == interval.lower + assert result.aime_lower == interval.lower + assert result.aime_upper == interval.upper + assert compute_pia(result.aime, 2020, bend_points=BENDS) == D(pia) + + +def test_nonterminating_lower_endpoint_is_not_rounded_down_before_forward_pia(): + result = compute_pia_ratio( + "1.2", 2020, REVIEW_PATH, bend_points=BENDS, exact_eligibility_pia=True + ) + assert result.factor == 1 + assert Fraction(result.aime) >= Fraction(1, 9) + + +def test_annual_measurement_flag_even_when_derived_pia_is_a_dime(): + result = compute_pia_ratio( + "18117.12", 2020, bend_points=BENDS, claiming_adjustment=".7" + ) + assert result.aime == 5000 + assert result.measurement_uncertainty + + +def test_annual_rounding_interval_is_propagated_through_all_divisors(): + point = compute_pia_ratio( + "18117.12", + 2020, + bend_points=BENDS, + claiming_adjustment=".7", + auxiliary_share=".5", + cola_factor=2, + ) + expanded = compute_pia_ratio( + "18117.12", + 2020, + bend_points=BENDS, + claiming_adjustment=".7", + auxiliary_share=".5", + cola_factor=2, + annual_rounding_bound=".5", + ) + assert expanded.aime == point.aime == 5000 + assert expanded.aime_lower < point.aime_lower + assert expanded.aime_upper > point.aime_upper + assert expanded.measurement_uncertainty + + +def test_survey_non_dime_interval_does_not_claim_to_be_exact(): + result = compute_pia_ratio( + "18117.13", 2020, bend_points=BENDS, claiming_adjustment=".7" + ) + assert result.measurement_uncertainty + assert result.aime == result.aime_lower > 5000 + assert result.aime_upper > result.aime + with pytest.raises(ValueError, match="multiple of one dime"): + compute_pia_ratio( + "18117.13", + 2020, + bend_points=BENDS, + claiming_adjustment=".7", + exact_eligibility_pia=True, + ) + + +def test_baseline_inverse_never_uses_reform_factors(): + lower_path = FormulaPath(factor_values=(".5", ".2", ".1")) + for path in (REVIEW_PATH, lower_path): + result = compute_pia_ratio( + "25881.6", 2020, path, bend_points=BENDS, exact_eligibility_pia=True + ) + assert result.aime == 5000 + + +def test_custom_unreformed_historical_scale_is_used_for_both_inverse_and_pia(): + result = compute_pia_ratio( + 1200, + 2020, + FormulaPath(factor_values=(".9", ".32", ".15")), + bend_points=BENDS, + factors=(".8", ".3", ".1"), + exact_eligibility_pia=True, + ) + assert result.aime == 125 + assert result.factor == D("1.125") + + +def test_colas_are_removed_before_inversion_once(): + result = compute_pia_ratio( + "36234.24", + 2020, + REVIEW_PATH, + bend_points=BENDS, + claiming_adjustment=".7", + cola_factor=2, + exact_eligibility_pia=True, + ) + assert result.aime == 5000 + assert D("36234.24") * result.factor == D("38564.4") + + +def test_zero_benefit_returns_neutral_without_inversion_or_formula_inputs(): + result = compute_pia_ratio(0, 2020, REVIEW_PATH, bend_points=(0, 0)) + assert result.factor == 1 + assert result.aime == result.aime_lower == result.aime_upper == 0 + assert D(0) * result.factor == 0 + assert not result.measurement_uncertainty + + +def test_pre1983_record_is_flagged_and_left_unreformed_without_inversion(): + result = compute_pia_ratio(20000, 1982, REVIEW_PATH, bend_points=(0, 0)) + assert result.factor == 1 + assert result.pre1983 + assert result.aime == result.aime_lower == result.aime_upper == 0 + + +@pytest.mark.parametrize("benefit", [-1, "NaN", "Infinity", True]) +def test_invalid_reported_benefit_is_rejected_even_before_1983(benefit): + with pytest.raises(ValueError, match="annual reported benefit"): + compute_pia_ratio(benefit, 1982, bend_points=BENDS) + + +@pytest.mark.parametrize( + "name", ["claiming_adjustment", "auxiliary_share", "cola_factor"] +) +@pytest.mark.parametrize("value", [0, -1, "NaN", "Infinity"]) +def test_invalid_divisor_is_rejected(name, value): + with pytest.raises(ValueError, match="positive|finite"): + compute_pia_ratio(10000, 2020, bend_points=BENDS, **{name: value}) + + +def test_nonpositive_baseline_pia_is_not_an_infinite_factor(): + with pytest.raises(ValueError, match="baseline PIA must be positive"): + compute_pia_ratio(".01", 2020, REVIEW_PATH, bend_points=BENDS) + + +@pytest.mark.parametrize("bound", [-1, "NaN", "Infinity"]) +def test_invalid_measurement_bound_rejected(bound): + with pytest.raises(ValueError, match="annual rounding bound"): + compute_pia_ratio(10000, 2020, bend_points=BENDS, annual_rounding_bound=bound) + + +def test_exact_pia_cannot_also_have_an_annual_measurement_interval(): + with pytest.raises(ValueError, match="cannot have an annual rounding bound"): + compute_pia_ratio( + 1200, + 2020, + bend_points=BENDS, + exact_eligibility_pia=True, + annual_rounding_bound=1, + ) + + +def test_cohort_and_shared_retroactive_modes_have_distinct_applicability(): + path = FormulaPath(factor_values=(".95", ".32", ".15"), factor_start_year=2027) + assert compute_pia_ratio("25881.6", 2020, path, 2027, bend_points=BENDS).factor == 1 + assert ( + compute_pia_ratio( + "25881.6", 2020, path, 2026, "retroactive", bend_points=BENDS + ).factor + == 1 + ) + assert ( + compute_pia_ratio( + "25881.6", 2020, path, 2027, "retroactive", bend_points=BENDS + ).factor + > 1 + ) + + +@given( + st.decimals(min_value=".01", max_value="1000000", places=2), st.integers(1983, 2035) +) +def test_every_reform_off_is_exactly_one(benefit, year): + assert compute_pia_ratio(benefit, year, bend_points=BENDS).factor == D(1) + + +@given(st.integers(1, 50000)) +def test_baseline_inverse_roundtrips_and_contains_the_original_aime(aime): + pia = compute_pia(aime, 2020, bend_points=BENDS) + result = compute_pia_ratio( + 12 * pia, 2020, bend_points=BENDS, exact_eligibility_pia=True + ) + assert result.aime_lower <= aime < result.aime_upper + assert compute_pia(result.aime, 2020, bend_points=BENDS) == pia + + +@given(st.integers(1, 50000)) +def test_everywhere_higher_scale_has_factor_at_least_one(aime): + pia = compute_pia(aime, 2020, bend_points=BENDS) + result = compute_pia_ratio( + 12 * pia, 2020, REVIEW_PATH, bend_points=BENDS, exact_eligibility_pia=True + ) + assert result.factor >= 1 + + +@given(st.integers(1, 50000), st.integers(0, 100)) +def test_factor_is_never_negative_for_nonnegative_reform_scale(aime, percent): + pia = compute_pia(aime, 2020, bend_points=BENDS) + path = FormulaPath(factor_values=(Fraction(percent, 100), ".32", ".15")) + result = compute_pia_ratio( + 12 * pia, 2020, path, bend_points=BENDS, exact_eligibility_pia=True + ) + assert result.factor >= 0 + + +@given(st.integers(1983, 2035)) +def test_zero_remains_zero_under_an_active_reform(year): + result = compute_pia_ratio(0, year, REVIEW_PATH, bend_points=BENDS) + assert result.factor == 1 + assert 0 * result.factor == 0 diff --git a/policyengine_us/tools/social_security/pia_ratio.py b/policyengine_us/tools/social_security/pia_ratio.py new file mode 100644 index 00000000000..d4c704e6d47 --- /dev/null +++ b/policyengine_us/tools/social_security/pia_ratio.py @@ -0,0 +1,208 @@ +"""Infer an eligibility-cohort AIME from an annual reported benefit. + +This is the population, proportional-change approximation. It does not recover +an unidentified monthly payment history or undo earnings-test withholding, +dual entitlement or a binding family maximum. Callers must identify the worker, +claiming adjustment and auxiliary share before using it. Both the reform and +baseline PIAs use the shared household calculator's formula and dime rounding. +The inverse always uses the supplied *unreformed* historical scale. +""" + +from dataclasses import dataclass +from decimal import Decimal, ROUND_CEILING, localcontext +from fractions import Fraction + +from policyengine_us.tools.social_security.pia import ( + DEFAULT_FACTORS, + FormulaPath, + Number, + compute_pia, + has_active_reform, + inverse_pia, + inverse_unrounded, +) + + +@dataclass(frozen=True) +class RatioResult: + """A ratio and monthly-dollar inference diagnostics for one annual record. + + ``aime`` is the baseline inverse at the observed eligibility PIA. For an + exact eligibility PIA it is the lower endpoint of the half-open interval + ``[aime_lower, aime_upper)``. Survey observations instead give a conditional + interval, with additional measurement error explicitly flagged. Endpoints + are represented as 50-digit Decimals rounded toward positive infinity; + exact rational endpoints are retained internally for forward computation. + Zero and unsupported pre-1983 records have factor 1 and zero diagnostics. + """ + + factor: Decimal + aime: Decimal + aime_lower: Decimal + aime_upper: Decimal + pre1983: bool = False + measurement_uncertainty: bool = False + + +def _number(value: Number, name: str) -> Fraction: + try: + if isinstance(value, bool): + raise ValueError + if isinstance(value, Fraction): + return value + result = Decimal(str(value)) + if not result.is_finite(): + raise ValueError + return Fraction(result) + except (ValueError, TypeError, ArithmeticError) as error: + raise ValueError(f"{name} must be a finite number") from error + + +def _decimal(value: Fraction, *, ceiling: bool = False) -> Decimal: + with localcontext() as context: + context.prec = 50 + if ceiling: + context.rounding = ROUND_CEILING + return Decimal(value.numerator) / Decimal(value.denominator) + + +def compute_pia_ratio( + annual_benefit: Number, + eligibility_year: int, + formula_path: FormulaPath | None = None, + payment_year: int | None = None, + recompute_mode: str = "cohort", + *, + bend_points: tuple[Number, Number], + claiming_adjustment: Number = 1, + auxiliary_share: Number = 1, + cola_factor: Number = 1, + factors: tuple[Number, Number, Number] = DEFAULT_FACTORS, + exact_eligibility_pia: bool = False, + annual_rounding_bound: Number = 0, +) -> RatioResult: + """Return ``P_new(A, E; t) / P_old(A, E)`` for a reported annual benefit. + + The inferred monthly eligibility PIA is + ``annual_benefit / (12 * claiming_adjustment * auxiliary_share * cola_factor)``. + ``cola_factor`` must be the unreformed accumulation from eligibility, and + ``annual_benefit`` must be the unreformed reported amount. A supplied + composite survivor ratio belongs in either adjustment argument, once. + + Set ``exact_eligibility_pia=True`` only when the backwards-derived amount + is an exact dime-rounded eligibility PIA. Otherwise the point estimate is + the unrounded baseline inverse, and the interval extends through the + inverse at ``z + .10``. It is conditional on the observed survey amount; + ``measurement_uncertainty`` remains true even if z happens to be a dime. + An optional nonnegative ``annual_rounding_bound`` expands the annual + observation to ``[max(0, b-bound), b+bound)`` before dividing by 12*a*s*C. + The upper end additionally allows one dime of initial PIA truncation. + This conservative interval does not represent other survey measurement + errors or unidentified monthly payment rounding. It cannot be combined + with the exact-PIA assertion. + + Pre-1983 benchmark records remain unchanged because the historical + upward-rounding method is unsupported. Negative/nonfinite benefits and + nonpositive/nonfinite adjustment, share or COLA inputs are rejected. + """ + benefit = _number(annual_benefit, "annual reported benefit") + if benefit < 0: + raise ValueError("annual reported benefit must be nonnegative") + adjustments = tuple( + _number(value, name) + for value, name in ( + (claiming_adjustment, "claiming adjustment"), + (auxiliary_share, "auxiliary share"), + (cola_factor, "baseline COLA factor"), + ) + ) + if any(value <= 0 for value in adjustments): + raise ValueError( + "claiming adjustment, auxiliary share and COLA factor must be positive" + ) + bound = _number(annual_rounding_bound, "annual rounding bound") + if bound < 0: + raise ValueError("annual rounding bound must be nonnegative") + if exact_eligibility_pia and bound: + raise ValueError("exact eligibility PIA cannot have an annual rounding bound") + if isinstance(eligibility_year, bool) or not isinstance(eligibility_year, int): + raise ValueError("eligibility year must be an integer calendar year") + neutral = RatioResult(Decimal(1), Decimal(0), Decimal(0), Decimal(0)) + if eligibility_year < 1983: + return RatioResult( + neutral.factor, + neutral.aime, + neutral.aime_lower, + neutral.aime_upper, + pre1983=True, + ) + if benefit == 0: + return neutral + + divisor = 12 * adjustments[0] * adjustments[1] * adjustments[2] + pia = benefit / divisor + if exact_eligibility_pia or (pia * 10).denominator == 1: + interval = inverse_pia( + pia, eligibility_year, bend_points=bend_points, factors=factors + ) + point = interval.exact_lower + lower, upper = interval.exact_lower, interval.exact_upper + else: + point = inverse_unrounded( + pia, eligibility_year, bend_points=bend_points, factors=factors + ) + lower = point + upper = inverse_unrounded( + pia + Fraction(1, 10), + eligibility_year, + bend_points=bend_points, + factors=factors, + ) + if bound: + lower = inverse_unrounded( + max(0, benefit - bound) / divisor, + eligibility_year, + bend_points=bend_points, + factors=factors, + ) + upper = inverse_unrounded( + (benefit + bound) / divisor + Fraction(1, 10), + eligibility_year, + bend_points=bend_points, + factors=factors, + ) + + active = has_active_reform( + formula_path, eligibility_year, payment_year, recompute_mode + ) + old = compute_pia( + point, + eligibility_year, + payment_year=payment_year, + recompute_mode=recompute_mode, + bend_points=bend_points, + factors=factors, + round_result=True, + ) + new = compute_pia( + point, + eligibility_year, + formula_path, + payment_year, + recompute_mode, + bend_points=bend_points, + factors=factors, + round_result=True, + ) + if (active or old != new) and old <= 0: + raise ValueError("inferred baseline PIA must be positive for a PIA ratio") + factor = Decimal(1) + if old != new: + factor = _decimal(Fraction(new) / Fraction(old)) + return RatioResult( + factor=factor, + aime=_decimal(point, ceiling=True), + aime_lower=_decimal(lower, ceiling=True), + aime_upper=_decimal(upper, ceiling=True), + measurement_uncertainty=not exact_eligibility_pia, + ) From cb4c658a1c0ae8716eccb9ceb993cfdd9ae5ea4c Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 27 Oct 2026 05:29:43 -0700 Subject: [PATCH 66/80] Share explicit worker computation bases across benefit levers --- .../tests/unit/test_ss_worker_basis.py | 136 ++++++++++++++++++ .../tools/social_security/records.py | 52 +++++++ .../tools/social_security_benefit_levers.py | 31 ++-- 3 files changed, 197 insertions(+), 22 deletions(-) create mode 100644 policyengine_us/tests/unit/test_ss_worker_basis.py create mode 100644 policyengine_us/tools/social_security/records.py diff --git a/policyengine_us/tests/unit/test_ss_worker_basis.py b/policyengine_us/tests/unit/test_ss_worker_basis.py new file mode 100644 index 00000000000..b87629fd0f9 --- /dev/null +++ b/policyengine_us/tests/unit/test_ss_worker_basis.py @@ -0,0 +1,136 @@ +"""Worker-link resolution without simulation or country-model fixtures.""" + +import numpy as np +import pytest + +from policyengine_us.tools.social_security.records import resolve_worker_basis + + +class Records: + ids = ("recipient", "worker") + + def __init__(self, component, **inputs): + self.inputs = { + f"social_security_{component}_worker_eligibility_year": np.array([0, 0]), + f"social_security_{component}_worker_id": np.array(["", ""], dtype=object), + "social_security_worker_computation_year": np.array([0, 0]), + "social_security_retirement_worker_eligibility_year": np.array([0, 0]), + } + self.inputs.update({key: np.asarray(value) for key, value in inputs.items()}) + + def __call__(self, variable, period): + return self.inputs[variable] + + +@pytest.mark.parametrize("component", ["retirement", "disability"]) +def test_unlinked_own_record_uses_shared_year_without_mutating_inputs(component): + key = f"social_security_{component}_worker_eligibility_year" + records = Records(component, social_security_worker_computation_year=[2020, 2021]) + years, linked, indices = resolve_worker_basis( + records, 2026, component, [True, True] + ) + np.testing.assert_array_equal(years, [2020, 2021]) + np.testing.assert_array_equal(linked, [False, False]) + np.testing.assert_array_equal(indices, [0, 1]) + np.testing.assert_array_equal(records.inputs[key], [0, 0]) + + +def test_supplied_component_year_precedes_own_shared_year(): + records = Records( + "retirement", + social_security_retirement_worker_eligibility_year=[2018, 0], + social_security_worker_computation_year=[2020, 2021], + ) + years, _, _ = resolve_worker_basis(records, 2026, "retirement", [True, True]) + np.testing.assert_array_equal(years, [2018, 2021]) + + +def test_link_reads_workers_same_component_year_before_shared_year(): + records = Records( + "survivors", + social_security_survivors_worker_id=["worker", ""], + social_security_survivors_worker_eligibility_year=[0, 2018], + social_security_worker_computation_year=[0, 2020], + ) + years, linked, indices = resolve_worker_basis( + records, 2026, "survivors", [True, False] + ) + assert years[0] == 2018 + np.testing.assert_array_equal(linked, [True, False]) + np.testing.assert_array_equal(indices, [1, 1]) + + +@pytest.mark.parametrize("component", ["survivors", "dependents"]) +def test_auxiliary_link_can_use_workers_retirement_year(component): + records = Records( + component, + **{ + f"social_security_{component}_worker_id": ["worker", ""], + "social_security_retirement_worker_eligibility_year": [0, 2019], + }, + ) + years, linked, _ = resolve_worker_basis(records, 2026, component, [True, False]) + assert years[0] == 2019 + assert linked[0] + + +def test_auxiliary_without_worker_link_does_not_borrow_own_shared_year(): + records = Records("survivors", social_security_worker_computation_year=[2020, 2021]) + years, linked, _ = resolve_worker_basis(records, 2026, "survivors", [True, True]) + np.testing.assert_array_equal(years, [0, 0]) + assert not np.any(linked) + + +def test_link_without_worker_year_stays_missing_for_callers_to_handle(): + records = Records("retirement", social_security_retirement_worker_id=["worker", ""]) + years, linked, _ = resolve_worker_basis(records, 2026, "retirement", [True, False]) + assert years[0] == 0 + assert linked[0] + + +def test_links_read_supplied_worker_basis_without_order_dependent_link_chasing(): + records = Records( + "survivors", + social_security_survivors_worker_id=["worker", "recipient", ""], + social_security_survivors_worker_eligibility_year=[0, 0, 2019], + social_security_worker_computation_year=[0, 0, 0], + social_security_retirement_worker_eligibility_year=[0, 0, 0], + ) + records.ids = ("recipient", "linked", "worker") + years, linked, _ = resolve_worker_basis( + records, 2026, "survivors", [True, True, False] + ) + np.testing.assert_array_equal(years, [2019, 0, 2019]) + np.testing.assert_array_equal(linked, [True, True, False]) + + +def test_recipient_supplied_year_can_identify_linked_worker_basis(): + records = Records( + "survivors", + social_security_survivors_worker_id=["worker", ""], + social_security_survivors_worker_eligibility_year=[2018, 0], + ) + years, linked, _ = resolve_worker_basis(records, 2026, "survivors", [True, False]) + assert years[0] == 2018 + assert linked[0] + + +def test_conflicting_supplied_and_linked_years_rejected(): + records = Records( + "survivors", + social_security_survivors_worker_id=["worker", ""], + social_security_survivors_worker_eligibility_year=[2018, 2019], + ) + with pytest.raises(ValueError, match="Conflicting"): + resolve_worker_basis(records, 2026, "survivors", [True, False]) + + +def test_unknown_worker_id_rejected_only_for_active_records(): + records = Records( + "retirement", social_security_retirement_worker_id=["missing", ""] + ) + with pytest.raises(ValueError, match="does not identify a person"): + resolve_worker_basis(records, 2026, "retirement", [True, False]) + years, linked, _ = resolve_worker_basis(records, 2026, "retirement", [False, True]) + assert years[0] == 0 + assert not np.any(linked) diff --git a/policyengine_us/tools/social_security/records.py b/policyengine_us/tools/social_security/records.py new file mode 100644 index 00000000000..563d01322ad --- /dev/null +++ b/policyengine_us/tools/social_security/records.py @@ -0,0 +1,52 @@ +"""Resolve explicitly supplied worker computation bases without age proxies.""" + +import numpy as np + + +def resolve_worker_basis(person, period, component, active): + """Return computation years, link flags and worker indices for each person. + + A recipient's component-specific computation year takes precedence over a + linked worker's supplied basis. If both are present they must agree. Links + read the worker's same-component year, then the shared computation-year + input; auxiliary links may also use the worker's retirement-specific year. + Unlinked own retirement/disability records may use their own shared year. + + Missing years remain zero. In particular, no recipient-age proxy is applied + here. Invalid links and conflicting years are rejected only for ``active`` + records, so an unused source cannot disturb a selected calculation route. + The returned arrays are private copies and never mutate stored inputs. + """ + prefix = f"social_security_{component}" + supplied_years = np.asarray( + person(f"{prefix}_worker_eligibility_year", period), dtype=int + ) + years = supplied_years.copy() + shared = person("social_security_worker_computation_year", period) + identifiers = person(f"{prefix}_worker_id", period) + known = {str(identifier): i for i, identifier in enumerate(person.ids)} + linked = np.zeros(len(years), dtype=bool) + worker_indices = np.arange(len(years), dtype=int) + auxiliary_years = ( + person("social_security_retirement_worker_eligibility_year", period) + if component in ("survivors", "dependents") + else None + ) + for i in np.flatnonzero(active): + identifier = str(identifiers[i]) + if not identifier: + if component in ("retirement", "disability"): + years[i] = years[i] or shared[i] + continue + if identifier not in known: + raise ValueError("Social Security worker_id does not identify a person") + j = known[identifier] + worker_indices[i] = j + record_year = supplied_years[j] or shared[j] + if not record_year and auxiliary_years is not None: + record_year = auxiliary_years[j] + if years[i] and record_year and years[i] != record_year: + raise ValueError("Conflicting supplied and linked worker computation years") + years[i] = years[i] or record_year + linked[i] = True + return years, linked, worker_indices diff --git a/policyengine_us/tools/social_security_benefit_levers.py b/policyengine_us/tools/social_security_benefit_levers.py index fb86fa14647..fc2484063f9 100644 --- a/policyengine_us/tools/social_security_benefit_levers.py +++ b/policyengine_us/tools/social_security_benefit_levers.py @@ -9,6 +9,7 @@ import numpy as np from policyengine_core.model_api import Enum +from policyengine_us.tools.social_security.records import resolve_worker_basis from policyengine_us.tools.ssa_cola_history import validate_cola_rounding_rule @@ -161,29 +162,11 @@ def reported_joint_cola_factor(person, period, parameters, component): 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") + worker_year, linked, _ = resolve_worker_basis(person, period, component, reported) + if np.any(reported & linked & (worker_year == 0) & (first_year <= 0)): + raise ValueError("Linked worker requires an explicit eligibility year") entitlement = person(f"{prefix}_entitlement_year", period) age = person("age", period) if component == "retirement": @@ -209,7 +192,11 @@ def reported_joint_cola_factor(person, period, parameters, component): if component == "retirement" else np.where(entitlement > 0, historical_first, proxy) ) - inferred = np.where(linked > 0, linked + (linked >= 1983), recipient_origin) + inferred = np.where( + worker_year > 0, + worker_year + (worker_year >= 1983), + recipient_origin, + ) first_year = np.where(first_year > 0, first_year, inferred) result = np.ones(len(age)) for event in events: From 86320a4423926f6f8ab5fbeab7d54b3f75f11502 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 27 Oct 2026 05:43:20 -0700 Subject: [PATCH 67/80] Connect reported Social Security PIA ratios and population diagnostics --- Makefile | 13 +- changelog.d/ss-pia-ratio-lever-c.added.md | 2 + docs/usage/social-security-pia-ratio.md | 217 +++++++++++ .../pia_ratio/annual_rounding_bound.yaml | 8 + .../pia_ratio/claiming_scenario.yaml | 7 + .../disability_eligibility_proxy.yaml | 7 + .../code_health/test_rest_python_groups.py | 57 ++- .../tests/core/test_ss_pia_ratio_inputs.py | 249 +++++++++++++ ..._security_dependents_pia_ratio_factor.yaml | 94 +++++ ..._security_disability_pia_ratio_factor.yaml | 60 +++ ..._security_retirement_pia_ratio_factor.yaml | 296 +++++++++++++++ ...l_security_survivors_pia_ratio_factor.yaml | 47 +++ .../unit/test_ss_pia_ratio_parameters.py | 57 +++ .../tests/unit/test_ss_population_survivor.py | 226 ++++++++++++ .../tools/social_security/population.py | 341 ++++++++++++++++++ .../social_security/population_survivor.py | 138 +++++++ .../gov/ssa/ss/social_security_birth_day.py | 10 + .../gov/ssa/ss/social_security_birth_month.py | 10 + ...ocial_security_dependents_inferred_aime.py | 16 + ...ty_dependents_inferred_aime_lower_bound.py | 18 + ...ty_dependents_inferred_aime_upper_bound.py | 18 + ...ependents_pia_ratio_claiming_adjustment.py | 16 + ...y_dependents_pia_ratio_eligibility_year.py | 16 + ...al_security_dependents_pia_ratio_factor.py | 9 +- ...dents_pia_ratio_measurement_uncertainty.py | 15 + ...pia_ratio_measurement_uncertainty_count.py | 22 ++ ..._dependents_pia_ratio_missing_claim_age.py | 15 + ...dents_pia_ratio_missing_claim_age_count.py | 22 ++ ...nts_pia_ratio_own_record_retired_worker.py | 15 + ..._security_dependents_pia_ratio_pre_1983.py | 15 + ...ity_dependents_pia_ratio_pre_1983_count.py | 20 + ...ial_security_dependents_pia_ratio_proxy.py | 15 + ...curity_dependents_pia_ratio_proxy_count.py | 20 + ...security_dependents_pia_ratio_unmodeled.py | 15 + ...ty_dependents_pia_ratio_unmodeled_count.py | 20 + ...ecurity_dependents_worker_benefit_ratio.py | 10 + ...rity_dependents_worker_ratio_identified.py | 10 + ...ocial_security_disability_inferred_aime.py | 16 + ...ty_disability_inferred_aime_lower_bound.py | 18 + ...ty_disability_inferred_aime_upper_bound.py | 18 + ...isability_pia_ratio_claiming_adjustment.py | 16 + ...y_disability_pia_ratio_eligibility_year.py | 16 + ...al_security_disability_pia_ratio_factor.py | 9 +- ...ility_pia_ratio_measurement_uncertainty.py | 15 + ...pia_ratio_measurement_uncertainty_count.py | 22 ++ ..._disability_pia_ratio_missing_claim_age.py | 15 + ...ility_pia_ratio_missing_claim_age_count.py | 22 ++ ...ity_pia_ratio_own_record_retired_worker.py | 15 + ..._security_disability_pia_ratio_pre_1983.py | 15 + ...ity_disability_pia_ratio_pre_1983_count.py | 20 + ...ial_security_disability_pia_ratio_proxy.py | 15 + ...curity_disability_pia_ratio_proxy_count.py | 20 + ...security_disability_pia_ratio_unmodeled.py | 15 + ...ty_disability_pia_ratio_unmodeled_count.py | 20 + ...al_security_retirement_claim_age_months.py | 10 + ...ocial_security_retirement_inferred_aime.py | 16 + ...ty_retirement_inferred_aime_lower_bound.py | 18 + ...ty_retirement_inferred_aime_upper_bound.py | 18 + ...etirement_pia_ratio_claiming_adjustment.py | 16 + ...y_retirement_pia_ratio_eligibility_year.py | 16 + ...al_security_retirement_pia_ratio_factor.py | 9 +- ...ement_pia_ratio_measurement_uncertainty.py | 15 + ...pia_ratio_measurement_uncertainty_count.py | 22 ++ ..._retirement_pia_ratio_missing_claim_age.py | 15 + ...ement_pia_ratio_missing_claim_age_count.py | 22 ++ ...ent_pia_ratio_own_record_retired_worker.py | 15 + ..._security_retirement_pia_ratio_pre_1983.py | 15 + ...ity_retirement_pia_ratio_pre_1983_count.py | 20 + ...ial_security_retirement_pia_ratio_proxy.py | 15 + ...curity_retirement_pia_ratio_proxy_count.py | 20 + ...security_retirement_pia_ratio_unmodeled.py | 15 + ...ty_retirement_pia_ratio_unmodeled_count.py | 20 + ...social_security_survivors_inferred_aime.py | 16 + ...ity_survivors_inferred_aime_lower_bound.py | 18 + ...ity_survivors_inferred_aime_upper_bound.py | 18 + ...survivors_pia_ratio_claiming_adjustment.py | 16 + ...ty_survivors_pia_ratio_eligibility_year.py | 16 + ...ial_security_survivors_pia_ratio_factor.py | 9 +- ...ivors_pia_ratio_measurement_uncertainty.py | 15 + ...pia_ratio_measurement_uncertainty_count.py | 22 ++ ...y_survivors_pia_ratio_missing_claim_age.py | 15 + ...ivors_pia_ratio_missing_claim_age_count.py | 22 ++ ...ors_pia_ratio_own_record_retired_worker.py | 15 + ...l_security_survivors_pia_ratio_pre_1983.py | 15 + ...rity_survivors_pia_ratio_pre_1983_count.py | 20 + ...cial_security_survivors_pia_ratio_proxy.py | 15 + ...ecurity_survivors_pia_ratio_proxy_count.py | 20 + ..._security_survivors_pia_ratio_unmodeled.py | 15 + ...ity_survivors_pia_ratio_unmodeled_count.py | 20 + ...security_survivors_worker_benefit_ratio.py | 10 + ...urity_survivors_worker_ratio_identified.py | 10 + 91 files changed, 3008 insertions(+), 29 deletions(-) create mode 100644 changelog.d/ss-pia-ratio-lever-c.added.md create mode 100644 docs/usage/social-security-pia-ratio.md create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/pia_ratio/annual_rounding_bound.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/pia_ratio/claiming_scenario.yaml create mode 100644 policyengine_us/parameters/gov/contrib/ssa/benefit_levers/pia_ratio/disability_eligibility_proxy.yaml create mode 100644 policyengine_us/tests/core/test_ss_pia_ratio_inputs.py create mode 100644 policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_dependents_pia_ratio_factor.yaml create mode 100644 policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_disability_pia_ratio_factor.yaml create mode 100644 policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml create mode 100644 policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_survivors_pia_ratio_factor.yaml create mode 100644 policyengine_us/tests/unit/test_ss_pia_ratio_parameters.py create mode 100644 policyengine_us/tests/unit/test_ss_population_survivor.py create mode 100644 policyengine_us/tools/social_security/population.py create mode 100644 policyengine_us/tools/social_security/population_survivor.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_birth_day.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_birth_month.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_inferred_aime.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_inferred_aime_lower_bound.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_inferred_aime_upper_bound.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_claiming_adjustment.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_eligibility_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_measurement_uncertainty.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_measurement_uncertainty_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_missing_claim_age.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_missing_claim_age_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_own_record_retired_worker.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_pre_1983.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_pre_1983_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_proxy.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_proxy_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_unmodeled.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_unmodeled_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_benefit_ratio.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_ratio_identified.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_inferred_aime.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_inferred_aime_lower_bound.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_inferred_aime_upper_bound.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_claiming_adjustment.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_eligibility_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_measurement_uncertainty.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_measurement_uncertainty_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_missing_claim_age.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_missing_claim_age_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_own_record_retired_worker.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_pre_1983.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_pre_1983_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_proxy.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_proxy_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_unmodeled.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_unmodeled_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_claim_age_months.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_inferred_aime.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_inferred_aime_lower_bound.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_inferred_aime_upper_bound.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_claiming_adjustment.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_eligibility_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_measurement_uncertainty.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_measurement_uncertainty_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_missing_claim_age.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_missing_claim_age_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_own_record_retired_worker.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_pre_1983.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_pre_1983_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_proxy.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_proxy_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_unmodeled.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_unmodeled_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_inferred_aime.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_inferred_aime_lower_bound.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_inferred_aime_upper_bound.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_claiming_adjustment.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_eligibility_year.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_measurement_uncertainty.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_measurement_uncertainty_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_missing_claim_age.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_missing_claim_age_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_own_record_retired_worker.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_pre_1983.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_pre_1983_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_proxy.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_proxy_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_unmodeled.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_unmodeled_count.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_benefit_ratio.py create mode 100644 policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_ratio_identified.py diff --git a/Makefile b/Makefile index c2efd4e5860..c4a8a7fd82c 100644 --- a/Makefile +++ b/Makefile @@ -83,6 +83,9 @@ REST_SS_CALCULATOR_DIFFERENTIAL_TESTS := $(TESTS)/unit/test_ss_calculator_differ # 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 +# Two-household H5 ingress and neutral-hook comparisons retain a reform tree. +# Isolate this file in the existing input target; no national dataset is loaded. +REST_SOCIAL_SECURITY_PIA_RATIO_INPUT_TESTS := $(TESTS)/core/test_ss_pia_ratio_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. @@ -98,7 +101,8 @@ REST_REMAINING_HEAVY_TESTS := $(TESTS)/test_niit_section_911_magi.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_SS_CALCULATOR_ENGINE_TESTS) $(REST_SS_CALCULATOR_DIFFERENTIAL_TESTS) \ - $(REST_SOCIAL_SECURITY_INPUT_TESTS) $(REST_POLICY_PARAMETER_TESTS) $(REST_REMAINING_HEAVY_TESTS) + $(REST_SOCIAL_SECURITY_INPUT_TESTS) $(REST_SOCIAL_SECURITY_PIA_RATIO_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 @@ -330,7 +334,14 @@ test-other-python-rest-ss-calculator-engine: test-other-python-rest-ss-calculator-differential: $(call rest_pytest,ss-calculator-differential,$(REST_SS_CALCULATOR_DIFFERENTIAL_TESTS)) test-other-python-rest-social-security-inputs: + @failed=""; for group in compatibility pia-ratio; do \ + $(MAKE) --no-print-directory test-other-python-rest-social-security-inputs-$$group || failed="$$failed $$group"; \ + done; \ + if [ -n "$$failed" ]; then echo "Failed Social Security input groups:$$failed"; exit 1; fi +test-other-python-rest-social-security-inputs-compatibility: $(call rest_pytest,social-security-inputs,$(REST_SOCIAL_SECURITY_INPUT_TESTS)) +test-other-python-rest-social-security-inputs-pia-ratio: + $(call rest_pytest,social-security-pia-ratio-inputs,$(REST_SOCIAL_SECURITY_PIA_RATIO_INPUT_TESTS)) test-other-python-rest-core: $(call rest_pytest,core,$(TESTS)/core,$(REST_LISTED_TESTS)) test-other-python-rest-policy-parameters: diff --git a/changelog.d/ss-pia-ratio-lever-c.added.md b/changelog.d/ss-pia-ratio-lever-c.added.md new file mode 100644 index 00000000000..501d39bcfb5 --- /dev/null +++ b/changelog.d/ss-pia-ratio-lever-c.added.md @@ -0,0 +1,2 @@ +Add an eligibility-year PIA-ratio lever for reported Social Security benefits, supplied claim-age inputs, and named population inference and coverage diagnostics. + diff --git a/docs/usage/social-security-pia-ratio.md b/docs/usage/social-security-pia-ratio.md new file mode 100644 index 00000000000..66cf4b43114 --- /dev/null +++ b/docs/usage/social-security-pia-ratio.md @@ -0,0 +1,217 @@ +# Population Social Security PIA ratio + +The reported-benefit route implements design v3 C1–C3, with A1's factor +composition, A5.2's record/proxy distinctions, P2's annual cross-section +convention, and amendments M2, M4, M8 and M9. It approximates a formula reform +for benefits already in payment. It does not reconstruct a worker's earnings +history or establish benefit entitlement. + +For each annual Person component `X` (`retirement`, `disability`, +`survivors`, `dependents`), the population lever applies + +```text +z = reported_X / (12 * a * s * baseline_COLA(E, t)) +factor_X = new_PIA(inferred_AIME, E; t) / baseline_PIA(inferred_AIME, E) +``` + +The public formula multiplies the reported base by this factor, the joint COLA +factor, the across-the-board factor and the payable factor. Calculated source +routes compute their PIA and COLAs directly and bypass these population +factors. An initial public-name input is migrated to its `_reported` twin, +including public-name dataset columns in the scoring year; a later explicit +public write is a final override. See the benefit-lever ingress documentation. + +The forward PIA calculations and inverse use the same service as the +[supplied-record calculator](social-security-calculator.md). +`bend_points_at` selects the original eligibility-year bend points; +`formula_path_at` supplies the combined reform and its `recompute_mode`. +C does not add a second retroactive-recomputation switch. The inverse always +uses the fixed unreformed formula and the unreformed reported amount. + +## Supplied claim age and explicit scenarios + +`social_security_retirement_claim_age_months` is a pure annual Person input: +age in whole months at first entitlement, so 62 years 3 months is 747. +It has no formula, uprating, or positive-amount selection rule. A projected H5 +may store this column like any other input, with a value for each scoring year. + +`gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario` explicitly selects: + +| Value | Claiming adjustment source | +|---|---| +| `supplied_claim_age` (default) | The supplied age in whole months. | +| `entitlement_dates` | Component entitlement year/month and supplied birth date. | +| `age_62` | Sensitivity assuming age 62. | +| `age_65` | Sensitivity assuming age 65. | +| `fra` | Sensitivity assuming cohort full retirement age. | +| `age_70` | Sensitivity assuming age 70. | + +Core represents these named selector parameters as singleton lists. For example, +set `claiming_scenario` to `[age_65]` and `disability_eligibility_proxy` +to `[current_year]` in a parameter reform; the adapter reads the one selected +name. Bare scalar strings are unsupported by Core's parameter schema. + +All scenarios preserve cohort FRA, early months and delayed-credit rates. +Under `supplied_claim_age`, a missing value, a value before 62, or a value +after 70 leaves the affected factor at one and raises +`_pia_ratio_missing_claim_age`. There is no silent change to another scenario. +The `entitlement_dates` scenario requires a known month and day of birth, +supplied as `social_security_birth_month` and `social_security_birth_day`, +or a complete `ss_calculator_birth_date`. Zero month/day means unavailable. +Exact dates preserve SSA attainment on the day before a birthday, including +January 1 births. + +Where only a birth year is available for an age-month scenario, the cohort +calculation uses the documented non-January-1 approximation and raises +`_pia_ratio_proxy`. A claiming scenario is substantive: in C2, replacing the +correct age-62 adjustment of 0.70 with 1 infers AIME $2,978 instead of $5,000 +and changes the annual gain from $1,165.248 to $1,664.64 before rounding, an +overstatement of about 43%. + +## Eligibility years and auxiliary identification + +Supplied `social_security_X_worker_eligibility_year`, +`social_security_X_worker_id` links and the underlying worker's +`social_security_worker_computation_year` precede proxies. A supplied +entitlement/onset year supplies the component's applicable basis where +appropriate. A link must identify an existing Person and a supplied computation +basis; conflicting record years are unsupported. + +For own retirement without a supplied basis, each scoring year's eligibility +proxy is `E = t - age_t + 62`. Rebuild it independently for every annual +cross-section. A row whose age stays 70 has E=2018 in 2026 and E=2027 in 2035; +the model does not carry an inferred cohort or claiming history between those +rows. Integer age also leaves birthday/within-year ambiguity, so compare adjacent +eligibility years when scoring sensitivities. + +Disability uses `a = 1` and a supplied onset, entitlement or worker computation +year when available. Missing onset is explicitly controlled by +`gov.contrib.ssa.benefit_levers.pia_ratio.disability_eligibility_proxy`: +`pre_reform` is the upper-exposure assumption, using a supported benchmark +before the active reform; `current_year` assumes newly entitled in the scoring +year. Both are labelled proxies, not observed onset histories. + +For survivors and dependents, a recipient age cannot identify the underlying +worker's PIA. A supplied worker link/year alone does not establish the share. +The identified auxiliary route also requires: + +- `social_security_X_worker_benefit_ratio`: the composite `a * s`, including + every applicable claiming adjustment and recipient share/reduction. +- `social_security_X_worker_ratio_identified`: an explicit assertion that this + is one identified worker PIA, with no unaccounted dual entitlement, binding + family cap or earnings-test withholding. + +For reported retirement, a record that also has positive reported survivor, +dependent or disability amounts remains unmodeled for this lever. The mixed +recipient amounts do not identify whether retirement is one full worker PIA, +dual entitlement or a partial-year conversion. + +The inversion divides by that composite ratio once. Without a supported +worker basis and identified ratio, the amount keeps factor one and raises +`_pia_ratio_unmodeled`. No recipient-age worker proxy is enabled by default. +For this composite route, the named claiming-adjustment output is one as a +decomposition convention: the entire supplied `a * s` is assigned to the +worker-benefit ratio. This does not establish that the underlying worker +claimed at FRA or identify `a` separately. +The restricted supplied-ratio survivor household calculator retains its own +more detailed record-basis and nonbinding-limit assertions. + +When S.770 is active, reported survivors additionally require the same +recipient-side `ss_calculator_*` computation-basis fields as the household +route: worker birth year and death year, basis provenance, prior-disability +declaration, ordinary non-WINDEX assertion, and the applicable exact dates or +year-only/same-year assertions. An authoritative computation year is required +for prior disability. The basis must agree with the supplied worker eligibility +or computation year. Missing identification leaves factor one and raises +`_pia_ratio_unmodeled`; contradictory supplied information is rejected. +For a pre-2026 death before age 62 with no prior disability eligibility, the +existing S.770 interpretation switch controls only the 18% first-bend-point +increase. The 95% first factor applies with either switch value. + +Supported benchmark years start in 1983. A positive record with an earlier +benchmark keeps factor one and raises `_pia_ratio_pre_1983`. The initial PIA +rounding rule before 1983 is outside this inverse's scope. COLAs accumulate +from eligibility, including pre-claim COLAs, using the unreformed statutory +ledger; they do not use benefit-income uprating or reformed COLAs. + +## Named annual outputs + +Replace `X` below with any of the four component names. Monetary outputs are +monthly dollars at the eligibility-year valuation basis, even though the +reported source benefit is annual. + +| Person variable suffix after `social_security_X` | Meaning | +|---|---| +| `_inferred_aime` | Baseline inverse of the observed amount; the lower-endpoint convention for an exact dime-rounded eligibility PIA. | +| `_inferred_aime_lower_bound` | Inclusive conditional AIME interval bound. | +| `_inferred_aime_upper_bound` | Exclusive conditional AIME interval bound. | +| `_pia_ratio_eligibility_year` | Underlying worker eligibility/computation year E. | +| `_pia_ratio_claiming_adjustment` | Imputed worker claiming adjustment a; one by decomposition convention on the auxiliary composite-ratio route. | +| `_pia_ratio_own_record_retired_worker` | Own-record retired-worker indicator. | +| `_pia_ratio_missing_claim_age` | Required claiming information is unusable. | +| `_pia_ratio_pre_1983` | Benchmark predates supported initial-PIA rounding. | +| `_pia_ratio_unmodeled` | This amount lacks a supported identified inversion. | +| `_pia_ratio_proxy` | An eligibility, birth or claiming approximation was used. | +| `_pia_ratio_measurement_uncertainty` | Annual reported measurement is not an exact eligibility PIA. | + +Zero monetary outputs on records without an inference are placeholders. +Inspect the flags before interpreting them as zero earnings. The own-record +indicator describes the modeling classification; it is not evidence of insured +status or an administrative validation of benefit type. It is independent of +whether claiming data are missing or the benchmark predates 1983. For an AIME +distribution or percentile calculation, also exclude `_pia_ratio_unmodeled`, +`_pia_ratio_missing_claim_age` and `_pia_ratio_pre_1983` records and inspect +the proxy and measurement-uncertainty flags. + +For exact dime-rounded baseline eligibility PIA z, the interval is +`[inverse(z), inverse(z + 0.10))`; the point estimate is its lower endpoint. +C2's $2,156.80 baseline PIA therefore gives +`[5000, 5000.3125)`. When the source is a rounded paid benefit or a survey +annual amount, that narrow interval is conditional on the assumed measurement +basis. It does not quantify all survey error. +`gov.contrib.ssa.benefit_levers.pia_ratio.annual_rounding_bound` optionally +supplies a nonnegative symmetric annual-dollar bound for +`[max(0, b - bound), b + bound)`, propagated backward through the baseline +divisor, with one dime of initial PIA truncation added at the upper endpoint. +The point estimate continues to invert the observed amount; it does not move +to the lower endpoint of this wider measurement interval. The default zero +adds no annual measurement-resolution assumption; it does not certify exact +data. + +Core stores monetary input arrays as float32. Before exact arithmetic, the +adapter uses each stored scalar's shortest float32 decimal representation. +This avoids expanding binary representation noise into an additional PIA +dime; it does not round every annual source to cents or recover precision +already lost when the input was stored. + +A zero benefit produces factor one and zero monetary outputs without division. +Negative reported amounts and nonpositive divisors are unsupported and rejected. +Factors cannot be negative. With every PIA reform off, factors are exactly one. + +## Weighted counts and scoring limits + +For any Person indicator `flag`, the authoritative weighted count is + +```text +sum(flag_i * person_weight_i) +``` + +The associated unmodeled annual amount is +`sum(flag_i * social_security_X_reported_i * person_weight_i)`. +Apply weights once. Report identified results, proxy sensitivities and unmodeled +amounts separately. + +Household variables `social_security_X_pia_ratio_FLAG_count` are also available +for `FLAG` equal to `missing_claim_age`, `pre_1983`, `unmodeled`, `proxy` +or `measurement_uncertainty`. They sum the corresponding Person indicators +within a household. A normal weighted household aggregation gives the weighted +count when members use the household weight. If an artifact supplies distinct +Person weights within a household, use the Person-indicator formula above. + +This approximation assumes the reported annual amount represents a full-year +benefit at the stated valuation basis; partial-year entitlement, survey error, +dual entitlement, withheld checks and binding caps can defeat unique inversion. +No package-specific parameter values or administrative percentile estimates are +encoded here. M8 allows progressive price-indexing thresholds from a supplied +bend-point fraction or dollar path; C's named outputs remain available for +population diagnostics and drift checks. diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/pia_ratio/annual_rounding_bound.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/pia_ratio/annual_rounding_bound.yaml new file mode 100644 index 00000000000..d229b02dd5a --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/pia_ratio/annual_rounding_bound.yaml @@ -0,0 +1,8 @@ +description: "Nonnegative optional symmetric annual dollar rounding or measurement-resolution bound. A supplied reported amount b represents the half-open interval [max(0, b - bound), b + bound) before backward propagation through the baseline claiming, worker-share and COLA divisor. The upper endpoint also allows one dime of initial PIA truncation. Zero leaves only the conditional PIA-rounding interval; it does not assert survey measurement is exact." +values: + 1983-01-01: 0 +metadata: + unit: currency-USD + label: Social Security population PIA-ratio annual rounding bound + economy: false + period: year diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/pia_ratio/claiming_scenario.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/pia_ratio/claiming_scenario.yaml new file mode 100644 index 00000000000..c3daa60b20a --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/pia_ratio/claiming_scenario.yaml @@ -0,0 +1,7 @@ +description: "Explicit claiming-adjustment source: supplied_claim_age, entitlement_dates, age_62, age_65, fra, or age_70. supplied_claim_age reads the pure claim-age-months input and leaves missing or unusable records unreformed; there is no automatic fallback." +values: + 1983-01-01: [supplied_claim_age] +metadata: + unit: str + label: Social Security population PIA-ratio claiming scenario + economy: false diff --git a/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/pia_ratio/disability_eligibility_proxy.yaml b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/pia_ratio/disability_eligibility_proxy.yaml new file mode 100644 index 00000000000..3e813710fc0 --- /dev/null +++ b/policyengine_us/parameters/gov/contrib/ssa/benefit_levers/pia_ratio/disability_eligibility_proxy.yaml @@ -0,0 +1,7 @@ +description: "Explicit missing-onset sensitivity: pre_reform uses an eligibility benchmark before the earliest active reform path, bounded by supported years; current_year treats the record as newly entitled in the scoring year. Supplied entitlement and worker computation years take precedence. Both remain labelled proxies." +values: + 1983-01-01: [pre_reform] +metadata: + unit: str + label: Social Security disability PIA-ratio eligibility proxy + economy: false 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 33b0c24b2f8..f9560c504d5 100644 --- a/policyengine_us/tests/code_health/test_rest_python_groups.py +++ b/policyengine_us/tests/code_health/test_rest_python_groups.py @@ -952,22 +952,25 @@ def stub_environment(tmp_path, outcomes): return env, calls -@pytest.mark.parametrize("failed_index", [None, 0, 1, 2]) +@pytest.mark.parametrize("failed_index", [None, 0, 1, 2, 3]) 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.""" + """Keep four isolated processes, including both synthetic input files.""" report_option = "--junitxml=simulation-microsimulation.xml" 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 + input_commands = dry_run(input_target) + input_groups = [ + groups[name] + for name in ("social-security-inputs", "social-security-pia-ratio-inputs") + ] + assert input_commands == [["pytest", *group["args"]] for group in input_groups] + commands = [*microsimulation_commands, *input_commands] + assert len(commands) == 4 reports = [] process_files = {} for index, command in enumerate(microsimulation_commands): @@ -976,22 +979,29 @@ def test_social_security_processes_keep_reports_and_run_after_failures( process_files[str(index)] = collected_files( [a for a in command[1:] if not a.startswith("--junitxml=")] ) - 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]) + input_files = [ + "policyengine_us/tests/core/test_social_security_inputs.py", + "policyengine_us/tests/core/test_ss_pia_ratio_inputs.py", + ] + for index, (command, file) in enumerate(zip(input_commands, input_files)): + process_files[f"inputs-{index}"] = collected_files(command[1:]) + assert process_files[f"inputs-{index}"] == [file] + reference = collected_files([str(TESTS / "microsimulation"), *input_files]) 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 + for name, group in zip( + ("social-security-inputs", "social-security-pia-ratio-inputs"), input_groups + ): + (input_report,) = [ + report for report in group["reports"] if report.endswith(".xml") + ] + assert input_report == f"{REPORT_DIR}/rest-python-{name}.xml" + reports.append(f"--junitxml={input_report}") + assert len(set(reports)) == 4 - outcomes = ["0"] * 3 + outcomes = ["0"] * 4 if failed_index is not None: outcomes[failed_index] = "1" env, calls = stub_environment(tmp_path, outcomes) @@ -1013,7 +1023,7 @@ def test_social_security_processes_keep_reports_and_run_after_failures( ) assert calls.read_text().splitlines() == ["|".join(c[1:]) for c in commands] assert (microsimulation_result.returncode != 0) == (failed_index in (0, 1)) - assert (input_result.returncode != 0) == (failed_index == 2) + assert (input_result.returncode != 0) == (failed_index in (2, 3)) def scenario_outcomes(scenario, count): @@ -1044,7 +1054,14 @@ def test_every_group_runs_and_any_failure_fails_the_target(tmp_path, groups, sce failed = [group for group, outcome in zip(groups, outcomes) if outcome != "0"] assert (result.returncode != 0) == bool(failed), result.stdout + result.stderr if failed: - assert f"Failed Rest Python groups: {' '.join(failed)}" in result.stdout + # Both input report groups run inside the same existing Make target. + failed_targets = dict.fromkeys( + "social-security-inputs" + if group == "social-security-pia-ratio-inputs" + else group + for group in failed + ) + assert f"Failed Rest Python groups: {' '.join(failed_targets)}" in result.stdout @pytest.mark.parametrize("scenario", ["one-fails-one-killed", "all-fail", "all-pass"]) diff --git a/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py b/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py new file mode 100644 index 00000000000..d2ce8de4e6e --- /dev/null +++ b/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py @@ -0,0 +1,249 @@ +"""PIA-ratio dataset ingress and compatibility with the neutral stacked hooks. + +The two-household H5 exercises a dataset column that a situation-only policy +test cannot cover. The differential replaces only C's four factor formulas +with A's original constant-one hooks: all other stacked-base calculations and +the public-column migration remain live. No population microsimulation runs. +""" + +from copy import deepcopy + +import numpy as np +import pytest +from policyengine_core.data import Dataset +from policyengine_core.reforms import Reform +from policyengine_core.simulations import Simulation as CoreSimulation + + +YEAR = 2026 +COMPONENTS = ( + "social_security_retirement", + "social_security_disability", + "social_security_survivors", + "social_security_dependents", +) +BENEFITS = (*COMPONENTS, "social_security") +DOWNSTREAM = ( + "taxable_social_security", + "tax_unit_taxable_social_security", + "adjusted_gross_income", + "income_tax", + "household_net_income", + "ssi_unearned_income", + "snap_unearned_income", +) +AMOUNTS = ((18_440.64, 0, 0, 0), (12_000, 6_000, 1_200, 300)) +GROUP_ENTITIES = ("household", "tax_unit", "spm_unit", "family", "marital_unit") + + +def _write_h5(path): + """Write scored-year public benefits and whole-month claim-age columns.""" + import h5py + + ids = np.array([1, 2]) + columns = { + "person_id": ids, + "age": [62, 70], + "social_security_retirement_claim_age_months": [747, 0], + "taxable_interest_income": [30_000, 10], + "state_fips": [48, 48], + } + for entity in GROUP_ENTITIES: + columns[f"person_{entity}_id"] = ids + columns[f"{entity}_id"] = ids + columns[f"{entity}_weight"] = [2.0, 3.0] + 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}/{YEAR}", data=values) + + +def _situation(): + people = {} + for index, person in enumerate(("retiree", "missing_claim_age")): + people[person] = { + "age": {YEAR: (62, 70)[index]}, + "social_security_retirement_claim_age_months": {YEAR: (747, 0)[index]}, + "taxable_interest_income": {YEAR: (30_000, 10)[index]}, + **{ + component: {YEAR: AMOUNTS[index][column]} + for column, component in enumerate(COMPONENTS) + }, + } + result = {"people": people} + for entity in GROUP_ENTITIES: + result[f"{entity}s"] = {person: {"members": [person]} for person in people} + for household in result["households"].values(): + household["state_code"] = {YEAR: "TX"} + return result + + +@pytest.fixture(scope="module") +def country_system(): + # Share the read-only model; reforms and simulations keep independent state. + from policyengine_us.system import system + + return system + + +def _neutral_stacked_base(system): + """Restore the original A hooks while retaining the complete stacked model.""" + from policyengine_us.model_api import Person, Variable, YEAR + + class NeutralPIARatioHooks(Reform): + def apply(self): + for component in COMPONENTS: + variable = type( + f"{component}_pia_ratio_factor", + (Variable,), + { + "value_type": float, + "entity": Person, + "definition_period": YEAR, + "formula": lambda person, period, parameters: 1, + }, + ) + self.update_variable(variable) + + return NeutralPIARatioHooks(system) + + +def _snapshot(simulation): + requested = { + name: np.asarray(simulation.calculate(name, YEAR)).copy() + for name in (*BENEFITS, *DOWNSTREAM) + } + # Check intermediates reached by these outputs, not just their final + # totals. New diagnostic-only values have no counterpart in A's hooks. + cached = {} + for population in simulation.populations.values(): + for name, holder in population._holders.items(): + array = holder.get_array(YEAR, simulation.branch_name) + if array is not None: + cached[name] = np.asarray(array).copy() + return requested, cached + + +@pytest.mark.parametrize("ingress", ("reported_situation", "public_h5")) +def test_all_off_matches_neutral_stacked_base(country_system, tmp_path, ingress): + """C adds no current-law change on either supported initial-ingress route.""" + if ingress == "public_h5": + path = tmp_path / "reported_benefits_2026.h5" + _write_h5(path) + inputs = {"dataset": Dataset.from_file(path)} + else: + inputs = {"situation": _situation()} + # Explicit reported names are the policy-sensitive route in a situation. + for person in inputs["situation"]["people"].values(): + for component in COMPONENTS: + person[f"{component}_reported"] = person.pop(component) + + results = [] + for system in (country_system, _neutral_stacked_base(country_system)): + simulation_inputs = ( + {"dataset": Dataset.from_file(path)} + if ingress == "public_h5" + else deepcopy(inputs) + ) + simulation = CoreSimulation( + **simulation_inputs, + tax_benefit_system=system, + default_input_period=YEAR, + default_calculation_period=YEAR, + ) + requested, cached = _snapshot(simulation) + for index, component in enumerate(COMPONENTS): + np.testing.assert_array_equal( + requested[component], + np.asarray([amounts[index] for amounts in AMOUNTS], dtype=np.float32), + ) + results.append((requested, cached)) + + current, reference = results + assert current[0].keys() == reference[0].keys() + for name in current[0]: + np.testing.assert_array_equal( + current[0][name], reference[0][name], err_msg=name + ) + assert reference[1].keys() <= current[1].keys() + for name in reference[1]: + np.testing.assert_array_equal( + current[1][name], reference[1][name], err_msg=name + ) + + +def test_supplied_claim_months_load_from_h5_and_drive_inverse(country_system, tmp_path): + path = tmp_path / "claim_months_2026.h5" + _write_h5(path) + reform = Reform.from_dict( + { + "gov.contrib.ssa.pia.in_effect": {YEAR: True}, + "gov.contrib.ssa.pia.factors.first": {YEAR: 0.95}, + "gov.contrib.ssa.pia.bend_point_multipliers.first": {YEAR: 1.15}, + "gov.ssa.social_security.pia.bend_points.first": {YEAR: 960}, + "gov.ssa.social_security.pia.bend_points.second": {YEAR: 5785}, + } + ) + simulation = CoreSimulation( + dataset=Dataset.from_file(path), + tax_benefit_system=country_system, + reform=reform, + default_input_period=YEAR, + default_calculation_period=YEAR, + ) + np.testing.assert_array_equal( + simulation.calculate("social_security_retirement_claim_age_months", YEAR), + [747, 0], + ) + np.testing.assert_array_equal( + simulation.calculate("social_security_retirement_claim_age_months", YEAR + 1), + [0, 0], + ) + adjustment = simulation.calculate( + "social_security_retirement_pia_ratio_claiming_adjustment", YEAR + ) + assert adjustment[0] == pytest.approx(0.7125) + aime = simulation.calculate("social_security_retirement_inferred_aime", YEAR) + assert aime[0] == pytest.approx(5000, abs=0.01) + lower = simulation.calculate( + "social_security_retirement_inferred_aime_lower_bound", YEAR + ) + upper = simulation.calculate( + "social_security_retirement_inferred_aime_upper_bound", YEAR + ) + assert lower[0] <= aime[0] < upper[0] + factors = simulation.calculate("social_security_retirement_pia_ratio_factor", YEAR) + assert factors[0] == pytest.approx(2295.5 / 2156.8) + assert factors[1] == 1 + np.testing.assert_array_equal( + simulation.calculate( + "social_security_retirement_pia_ratio_own_record_retired_worker", YEAR + ), + [True, False], + ) + flags = np.asarray( + simulation.calculate( + "social_security_retirement_pia_ratio_missing_claim_age", YEAR + ) + ) + np.testing.assert_array_equal(flags, [False, True]) + # Person flags retain the household weights supplied by the projected file. + weights = np.asarray(simulation.calculate("household_weight", YEAR)) + assert np.dot(flags, weights) == 3 + counts = simulation.calculate( + "social_security_retirement_pia_ratio_missing_claim_age_count", YEAR + ) + np.testing.assert_array_equal(counts, [0, 1]) + assert np.dot(counts, weights) == 3 + public = simulation.calculate("social_security_retirement", YEAR) + assert public[0] == pytest.approx(12 * 0.7125 * 2295.5, abs=0.01) + assert public[1] == 12_000 + + +def test_claim_months_are_a_pure_annual_input(country_system): + variable = country_system.variables["social_security_retirement_claim_age_months"] + assert variable.value_type is int + assert variable.default_value == 0 + assert not variable.formulas + assert not getattr(variable, "uprating", None) diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_dependents_pia_ratio_factor.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_dependents_pia_ratio_factor.yaml new file mode 100644 index 00000000000..6a61d790e13 --- /dev/null +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_dependents_pia_ratio_factor.yaml @@ -0,0 +1,94 @@ +- name: An auxiliary recipient age cannot identify the worker PIA + period: 2026 + input: + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + people: + child: + age: 10 + social_security_dependents_reported: 6000 + older_spouse: + age: 70 + social_security_dependents_reported: 6000 + output: + social_security_dependents_pia_ratio_factor: [1, 1] + social_security_dependents_pia_ratio_unmodeled: [true, true] + social_security_dependents_pia_ratio_own_record_retired_worker: [false, false] + social_security_dependents: [6000, 6000] + +- name: Linked dependent with an identified composite ratio uses the worker year + period: 2026 + absolute_error_margin: 0.01 + input: + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + gov.contrib.ssa.pia.bend_point_multipliers.first: 1.15 + people: + child: + age: 10 + social_security_dependents_worker_id: worker + social_security_dependents_worker_benefit_ratio: 0.5 + social_security_dependents_worker_ratio_identified: true + social_security_dependents_reported: 12940.8 + worker: + age: 62 + social_security_worker_computation_year: 2026 + output: + people: + child: + social_security_dependents_pia_ratio_eligibility_year: 2026 + social_security_dependents_pia_ratio_factor: 1.0643082344 + social_security_dependents_inferred_aime: 5000 + social_security_dependents_pia_ratio_unmodeled: false + social_security_dependents: 13773 + worker: + social_security_dependents_pia_ratio_factor: 1 + +- name: A worker link without an identified share remains unmodeled + period: 2026 + input: + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + people: + child: + age: 10 + social_security_dependents_worker_id: worker + social_security_dependents_reported: 6000 + worker: + age: 62 + social_security_worker_computation_year: 2026 + output: + people: + child: + social_security_dependents_pia_ratio_factor: 1 + social_security_dependents_pia_ratio_unmodeled: true + social_security_dependents: 6000 + +- name: A supplied 80 percent composite ratio preserves the exact baseline dime + period: 2026 + absolute_error_margin: 0.005 + input: + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + people: + dependent: + age: 70 + social_security_dependents_worker_id: worker + social_security_dependents_worker_benefit_ratio: 0.8 + social_security_dependents_worker_ratio_identified: true + # 12 * 0.8 * 2156.8. The input ratio is stored in a float32 holder. + social_security_dependents_reported: 20705.28 + worker: + age: 62 + social_security_worker_computation_year: 2026 + output: + people: + dependent: + social_security_dependents_inferred_aime: 5000 + social_security_dependents_pia_ratio_factor: 1.0222552 + social_security_dependents_pia_ratio_unmodeled: false + social_security_dependents: 21166.08 diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_disability_pia_ratio_factor.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_disability_pia_ratio_factor.yaml new file mode 100644 index 00000000000..adaa895567a --- /dev/null +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_disability_pia_ratio_factor.yaml @@ -0,0 +1,60 @@ +- name: A supplied onset year identifies disability PIA without a claiming reduction + period: 2026 + absolute_error_margin: 0.01 + input: + age: 45 + social_security_disability_entitlement_year: 2026 + social_security_disability_reported: 25881.60 + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + gov.contrib.ssa.pia.bend_point_multipliers.first: 1.15 + output: + social_security_disability_pia_ratio_eligibility_year: 2026 + social_security_disability_pia_ratio_claiming_adjustment: 1 + social_security_disability_pia_ratio_missing_claim_age: false + social_security_disability_inferred_aime: 5000 + social_security_disability_pia_ratio_factor: 1.0643082344 + social_security_disability: 27546 + social_security_disability_pia_ratio_proxy: false + social_security_disability_pia_ratio_own_record_retired_worker: false + +- name: Missing disability onset is an explicitly labelled pre reform proxy + period: 2026 + input: + age: 45 + social_security_disability_reported: 12000 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + gov.contrib.ssa.pia.factors.start_year: 2026 + output: + social_security_disability_pia_ratio_eligibility_year: 2025 + social_security_disability_pia_ratio_claiming_adjustment: 1 + social_security_disability_pia_ratio_proxy: true + social_security_disability_pia_ratio_factor: 1 + social_security_disability: 12000 + +- name: Current year disability proxy is selected by parameter + period: 2026 + input: + age: 45 + social_security_disability_reported: 12000 + gov.contrib.ssa.benefit_levers.pia_ratio.disability_eligibility_proxy: [current_year] + output: + social_security_disability_pia_ratio_eligibility_year: 2026 + social_security_disability_pia_ratio_proxy: true + social_security_disability_pia_ratio_factor: 1 + +- name: Supplied disability eligibility overrides the proxy choice + period: 2026 + input: + age: 45 + social_security_disability_worker_eligibility_year: 2018 + social_security_disability_entitlement_year: 2020 + social_security_disability_reported: 12000 + gov.contrib.ssa.benefit_levers.pia_ratio.disability_eligibility_proxy: [current_year] + output: + social_security_disability_pia_ratio_eligibility_year: 2018 + social_security_disability_pia_ratio_proxy: false + social_security_disability_pia_ratio_factor: 1 diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml new file mode 100644 index 00000000000..e3e181f3e3b --- /dev/null +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml @@ -0,0 +1,296 @@ +- name: Inactive population PIA reform preserves a reported record without ages + period: 2026 + input: + social_security_retirement_reported: 12000 + output: + social_security_retirement_pia_ratio_factor: 1 + social_security_retirement: 12000 + +- name: C2 lower-endpoint inverse preserves the 15 percent bend fixture + period: 2026 + absolute_error_margin: 0.01 + input: + age: 62 + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 18117.12 + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + gov.contrib.ssa.pia.bend_point_multipliers.first: 1.15 + output: + social_security_retirement_pia_ratio_eligibility_year: 2026 + social_security_retirement_pia_ratio_claiming_adjustment: 0.70 + social_security_retirement_inferred_aime: 5000 + social_security_retirement_inferred_aime_lower_bound: 5000 + social_security_retirement_inferred_aime_upper_bound: 5000.3125 + social_security_retirement_pia_ratio_factor: 1.0643082344 + social_security_retirement: 19282.20 + social_security_retirement_pia_ratio_own_record_retired_worker: true + social_security_retirement_pia_ratio_measurement_uncertainty: true + social_security_retirement_pia_ratio_unmodeled: false + +- name: FRA is an explicit sensitivity and changes inferred AIME + period: 2026 + absolute_error_margin: 0.01 + input: + age: 62 + social_security_retirement_reported: 18117.12 + gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [fra] + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + gov.contrib.ssa.pia.bend_point_multipliers.first: 1.15 + output: + social_security_retirement_pia_ratio_claiming_adjustment: 1 + social_security_retirement_inferred_aime: 2978 + social_security_retirement_pia_ratio_proxy: true + +- name: Age 62 sensitivity identifies its claiming adjustment explicitly + period: 2026 + absolute_error_margin: 0.01 + input: + age: 62 + social_security_retirement_reported: 18117.12 + gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [age_62] + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_claiming_adjustment: 0.70 + social_security_retirement_inferred_aime: 5000 + social_security_retirement_pia_ratio_proxy: true + +- name: Age 65 sensitivity retains the 24 month early retirement reduction + period: 2026 + absolute_error_margin: 0.01 + input: + age: 62 + social_security_retirement_reported: 22430.72 + gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [age_65] + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_claiming_adjustment: 0.8666666666666667 + social_security_retirement_inferred_aime: 5000 + social_security_retirement_pia_ratio_proxy: true + +- name: Age 70 sensitivity uses the birth cohort delayed retirement credits + period: 2026 + absolute_error_margin: 0.01 + input: + age: 62 + social_security_retirement_reported: 32093.184 + gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [age_70] + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_claiming_adjustment: 1.24 + social_security_retirement_inferred_aime: 5000 + social_security_retirement_pia_ratio_proxy: true + +- name: Missing or out of range supplied claim ages do not fall back + period: 2026 + input: + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + people: + missing: + age: 62 + social_security_retirement_reported: 12000 + before_62: + age: 62 + social_security_retirement_claim_age_months: 743 + social_security_retirement_reported: 12000 + after_70: + age: 72 + social_security_retirement_claim_age_months: 841 + social_security_retirement_reported: 12000 + output: + social_security_retirement_pia_ratio_factor: [1, 1, 1] + social_security_retirement_pia_ratio_missing_claim_age: [true, true, true] + social_security_retirement: [12000, 12000, 12000] + +- name: Zero benefit stays zero without a supplied claiming divisor + period: 2026 + input: + age: 62 + social_security_retirement_reported: 0 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_factor: 1 + social_security_retirement: 0 + social_security_retirement_inferred_aime: 0 + social_security_retirement_pia_ratio_missing_claim_age: false + +- name: Pre 1983 record remains unreformed and exposes its indicator + period: 2026 + input: + age: 106 + social_security_retirement_worker_eligibility_year: 1982 + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 12000 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.recompute_existing: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_eligibility_year: 1982 + social_security_retirement_pia_ratio_factor: 1 + social_security_retirement_pia_ratio_pre_1983: true + social_security_retirement_pia_ratio_pre_1983_count: 1 + social_security_retirement: 12000 + +- name: This scoring year rebuilds eligibility from its own age cross section + period: 2035 + input: + age: 70 + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 12000 + output: + social_security_retirement_pia_ratio_eligibility_year: 2027 + social_security_retirement_pia_ratio_factor: 1 + +- name: The same age in the 2026 cross section has a different eligibility year + period: 2026 + input: + age: 70 + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 12000 + output: + social_security_retirement_pia_ratio_eligibility_year: 2018 + social_security_retirement_pia_ratio_factor: 1 + +- name: Supplied entitlement and claim months override a conflicting age proxy + period: 2026 + absolute_error_margin: 0.0001 + input: + age: 70 + birth_year: 1960 + social_security_retirement_entitlement_year: 2025 + social_security_retirement_entitlement_month: 6 + social_security_retirement_claim_age_months: 780 + social_security_retirement_reported: 12000 + output: + # First entitlement at 65 in June 2025 identifies June 2022 eligibility. + social_security_retirement_pia_ratio_eligibility_year: 2022 + social_security_retirement_pia_ratio_claiming_adjustment: 0.8666666666666667 + +- name: A supplied birth year determines the full retirement age cohort + period: 2026 + absolute_error_margin: 0.01 + input: + age: 70 + birth_year: 1964 + social_security_retirement_worker_eligibility_year: 2026 + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 18117.12 + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + # The supplied cohort has FRA 67 despite this cross section's age proxy. + social_security_retirement_pia_ratio_eligibility_year: 2026 + social_security_retirement_pia_ratio_claiming_adjustment: 0.70 + social_security_retirement_inferred_aime: 5000 + +- name: A linked worker computation year overrides the recipient age proxy + period: 2026 + input: + people: + recipient: + age: 80 + social_security_retirement_worker_id: worker + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 12000 + worker: + age: 62 + social_security_worker_computation_year: 2020 + output: + people: + recipient: + social_security_retirement_pia_ratio_eligibility_year: 2020 + worker: + social_security_retirement_pia_ratio_factor: 1 + +- name: Entitlement dates select the same 62 years 3 months claiming adjustment + period: 2026 + absolute_error_margin: 0.0001 + input: + age: 62 + social_security_birth_month: 6 + social_security_birth_day: 15 + social_security_retirement_entitlement_year: 2026 + social_security_retirement_entitlement_month: 9 + social_security_retirement_reported: 18440.64 + gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [entitlement_dates] + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_eligibility_year: 2026 + social_security_retirement_pia_ratio_claiming_adjustment: 0.7125 + social_security_retirement_pia_ratio_missing_claim_age: false + +- name: Cohort gate leaves an existing beneficiary unchanged + period: 2026 + input: + age: 68 + social_security_retirement_worker_eligibility_year: 2020 + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 24000 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + gov.contrib.ssa.pia.factors.start_year: 2026 + output: + social_security_retirement_pia_ratio_factor: 1 + social_security_retirement: 24000 + +- name: The same factor change recomputes an old cohort only when selected + period: 2026 + absolute_error_margin: 0.01 + input: + age: 68 + social_security_retirement_worker_eligibility_year: 2020 + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 24000 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.recompute_existing: true + gov.contrib.ssa.pia.factors.first: 0.95 + gov.contrib.ssa.pia.factors.start_year: 2026 + output: + social_security_retirement_pia_ratio_eligibility_year: 2020 + social_security_retirement_pia_ratio_claiming_adjustment: 0.7166666666666667 + +- name: Calculated sources bypass the population inverse under an active reform + period: 2026 + input: + social_security_retirement_source: supplied_pia + social_security_retirement_base: 10000 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_factor: 1 + social_security_retirement_pia_ratio_missing_claim_age: false + social_security_retirement: 10000 + +- name: Introduced S770 retains its distinct 18 percent first bend increase + period: 2026 + absolute_error_margin: 0.01 + input: + age: 62 + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 18117.12 + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.congress.s770.section_2.in_effect: true + output: + social_security_retirement_inferred_aime: 5000 + social_security_retirement_pia_ratio_factor: 1.0727002967 + social_security_retirement: 19434.24 diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_survivors_pia_ratio_factor.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_survivors_pia_ratio_factor.yaml new file mode 100644 index 00000000000..a492f9bbbb4 --- /dev/null +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_survivors_pia_ratio_factor.yaml @@ -0,0 +1,47 @@ +- name: Survivor ages and recipient entitlement dates do not identify worker PIA + period: 2026 + input: + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + people: + young_survivor: + age: 10 + social_security_survivors_entitlement_year: 2025 + social_security_survivors_reported: 6000 + older_survivor: + age: 70 + social_security_survivors_reported: 12000 + output: + social_security_survivors_pia_ratio_factor: [1, 1] + social_security_survivors_pia_ratio_unmodeled: [true, true] + social_security_survivors: [6000, 12000] + +- name: A linked survivor composite ratio enters the inverse only once + period: 2026 + absolute_error_margin: 0.01 + input: + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + gov.contrib.ssa.pia.bend_point_multipliers.first: 1.15 + people: + survivor: + age: 60 + social_security_survivors_worker_id: worker + social_security_survivors_worker_benefit_ratio: 0.7 + social_security_survivors_worker_ratio_identified: true + social_security_survivors_reported: 18117.12 + worker: + age: 62 + social_security_worker_computation_year: 2026 + output: + people: + survivor: + social_security_survivors_pia_ratio_eligibility_year: 2026 + social_security_survivors_inferred_aime: 5000 + social_security_survivors_pia_ratio_factor: 1.0643082344 + social_security_survivors_pia_ratio_unmodeled: false + social_security_survivors: 19282.20 + worker: + social_security_survivors_pia_ratio_factor: 1 diff --git a/policyengine_us/tests/unit/test_ss_pia_ratio_parameters.py b/policyengine_us/tests/unit/test_ss_pia_ratio_parameters.py new file mode 100644 index 00000000000..cba3ae13975 --- /dev/null +++ b/policyengine_us/tests/unit/test_ss_pia_ratio_parameters.py @@ -0,0 +1,57 @@ +"""Load only C's authored parameter subtree, without a US model or simulation.""" + +from pathlib import Path + +import pytest +from policyengine_core.parameters import ParameterNode + + +PARAMETER_PATH = ( + Path(__file__).resolve().parents[2] + / "parameters/gov/contrib/ssa/benefit_levers/pia_ratio" +) +PARAMETER_NAME = "gov.contrib.ssa.benefit_levers.pia_ratio" + + +@pytest.fixture +def parameters(): + # The native Core loader must read the authored files. A hand-built dict or + # mocked snapshot would miss unsupported scalar-string parameter syntax. + return ParameterNode( + name=PARAMETER_NAME, + directory_path=str(PARAMETER_PATH), + ) + + +@pytest.mark.parametrize("instant", ["1983-01-01", "2026-01-01", "2035-01-01"]) +def test_authored_defaults_load_as_native_named_selector_lists(parameters, instant): + snapshot = parameters(instant) + assert snapshot.claiming_scenario == ["supplied_claim_age"] + assert snapshot.disability_eligibility_proxy == ["pre_reform"] + assert snapshot.annual_rounding_bound == 0 + + +def test_authored_selector_paths_and_metadata_survive_loading(parameters): + assert parameters.claiming_scenario.name == PARAMETER_NAME + ".claiming_scenario" + assert ( + parameters.disability_eligibility_proxy.name + == PARAMETER_NAME + ".disability_eligibility_proxy" + ) + assert parameters.claiming_scenario.metadata["unit"] == "str" + assert parameters.disability_eligibility_proxy.metadata["unit"] == "str" + + +def test_named_selector_reform_and_snapshots_keep_dated_values(parameters): + parameters.claiming_scenario.update(start="2026-01-01", value=["entitlement_dates"]) + parameters.disability_eligibility_proxy.update( + start="2026-01-01", value=["current_year"] + ) + parameters.annual_rounding_bound.update(start="2026-01-01", value=0.5) + earlier = parameters("2025-01-01") + current = parameters("2026-01-01") + assert earlier.claiming_scenario == ["supplied_claim_age"] + assert earlier.disability_eligibility_proxy == ["pre_reform"] + assert earlier.annual_rounding_bound == 0 + assert current.claiming_scenario == ["entitlement_dates"] + assert current.disability_eligibility_proxy == ["current_year"] + assert current.annual_rounding_bound == 0.5 diff --git a/policyengine_us/tests/unit/test_ss_population_survivor.py b/policyengine_us/tests/unit/test_ss_population_survivor.py new file mode 100644 index 00000000000..5ab05d52172 --- /dev/null +++ b/policyengine_us/tests/unit/test_ss_population_survivor.py @@ -0,0 +1,226 @@ +"""Survivor population scope uses the shared basis without building a model.""" + +from datetime import date +from decimal import Decimal +from types import SimpleNamespace + +import pytest + +from policyengine_us.tools.social_security.pia import FormulaPath, compute_pia +from policyengine_us.tools.social_security.population_survivor import ( + survivor_formula_path, +) + + +PERIOD = SimpleNamespace(start=date(2026, 1, 1)) +DEFAULTS = { + "worker_birth_year": 0, + "death_year": 0, + "authoritative_year": 0, + "prior_disability_eligibility": -1, + "same_year_death_before_62": -1, + "worker_birth_date": "", + "death_date": "", + "basis_provenance": "", + "ordinary_non_windex": False, + "year_only_not_january_1": False, +} +COMPLETE_BASIS = { + "worker_birth_year": 1970, + "death_year": 2025, + "prior_disability_eligibility": 0, + "basis_provenance": "Supplied ordinary worker computation basis", + "ordinary_non_windex": True, + "year_only_not_january_1": True, +} + + +class PersonInputs: + ids = ("survivor", "worker") + + def __init__(self, **values): + self.values = {**DEFAULTS, **values} + + def __call__(self, name, period): + assert period is PERIOD + field = name.removeprefix("ss_calculator_") + # The complete basis belongs to the recipient. The linked row has no + # assertions: the adapter must not combine partial bases across people. + return [self.values[field], DEFAULTS[field]] + + +def parameters(increase): + return SimpleNamespace( + gov=SimpleNamespace( + contrib=SimpleNamespace( + congress=SimpleNamespace( + s770=SimpleNamespace( + section_2=lambda when: SimpleNamespace( + recompute_pre2026_death_before_eligibility=increase + ) + ) + ) + ) + ) + ) + + +def resolve(values=None, *, increase=False, eligibility_year=2025, worker_index=1): + return survivor_formula_path( + PersonInputs(**(COMPLETE_BASIS if values is None else values)), + PERIOD, + parameters(increase), + 0, + FormulaPath(s770=True), + eligibility_year, + worker_index, + ) + + +def test_custom_reform_without_s770_does_not_require_a_deceased_worker_basis(): + path = FormulaPath(factor_values=(".95", ".32", ".15")) + actual, unsupported = survivor_formula_path(None, PERIOD, None, 0, path, 2025, None) + assert actual is path + assert not unsupported + + +@pytest.mark.parametrize( + ("increase", "gain_per_first_bend"), + [(False, Decimal(".05")), (True, Decimal(".1634"))], +) +def test_preeligibility_death_switch_changes_only_bend_increase( + increase, gain_per_first_bend +): + path, unsupported = resolve(increase=increase) + assert not unsupported + assert path.s770 + assert path.s770_increase_first_bend_point is increase + baseline = compute_pia(5000, 2025, bend_points=(960, 5785), round_result=False) + reform = compute_pia( + 5000, + 2025, + path, + 2026, + "retroactive", + bend_points=(960, 5785), + round_result=False, + ) + assert reform - baseline == gain_per_first_bend * 960 + + +def test_post2025_preeligibility_death_receives_both_increases_with_switch_off(): + path, unsupported = resolve( + {**COMPLETE_BASIS, "death_year": 2026}, eligibility_year=2026 + ) + assert not unsupported + assert path.s770 and path.s770_increase_first_bend_point + + +def test_death_after62_is_outside_the_preeligibility_switch(): + path, unsupported = resolve( + {**COMPLETE_BASIS, "worker_birth_year": 1950}, eligibility_year=2012 + ) + assert not unsupported + assert path.s770 and path.s770_increase_first_bend_point + + +def test_authoritative_prior_disability_basis_receives_both_increases(): + path, unsupported = resolve( + { + **COMPLETE_BASIS, + "prior_disability_eligibility": 1, + "authoritative_year": 2020, + }, + eligibility_year=2020, + ) + assert not unsupported + assert path.s770 and path.s770_increase_first_bend_point + + +def test_exact_january1_birth_uses_previous_attainment_year(): + path, unsupported = resolve( + { + **COMPLETE_BASIS, + "worker_birth_year": 1964, + "worker_birth_date": "1964-01-01", + "death_year": 2026, + "death_date": "2026-01-01", + "year_only_not_january_1": False, + }, + eligibility_year=2025, + ) + assert not unsupported + assert path.s770_increase_first_bend_point + + +@pytest.mark.parametrize( + "missing", + [ + {"worker_birth_year": 0}, + {"death_year": 0}, + {"prior_disability_eligibility": -1}, + {"basis_provenance": ""}, + {"ordinary_non_windex": False}, + {"year_only_not_january_1": False}, + {"prior_disability_eligibility": 1, "authoritative_year": 0}, + ], +) +def test_unavailable_basis_is_flagged_without_inference(missing): + path, unsupported = resolve({**COMPLETE_BASIS, **missing}) + assert unsupported + assert path.s770 and path.s770_increase_first_bend_point + + +def test_completely_absent_basis_is_unmodeled(): + _, unsupported = resolve({}) + assert unsupported + + +def test_sameyear_death_without_dates_or_assertion_is_unmodeled(): + _, unsupported = resolve({**COMPLETE_BASIS, "worker_birth_year": 1963}) + assert unsupported + + +@pytest.mark.parametrize(("before62", "increase"), [(0, True), (1, False)]) +def test_sameyear_explicit_assertion_resolves_the_shared_basis(before62, increase): + path, unsupported = resolve( + { + **COMPLETE_BASIS, + "worker_birth_year": 1963, + "same_year_death_before_62": before62, + } + ) + assert not unsupported + assert path.s770_increase_first_bend_point is increase + + +@pytest.mark.parametrize( + "invalid", + [ + {"worker_birth_year": -1}, + {"death_year": -1}, + {"authoritative_year": -1}, + {"prior_disability_eligibility": 2}, + {"same_year_death_before_62": 2}, + {"worker_birth_date": "invalid"}, + {"death_date": "2025-02-30"}, + {"worker_birth_date": "1971-01-02"}, + {"death_date": "2024-01-02"}, + {"worker_birth_year": 1990, "death_year": 1980}, + {"death_year": 2027}, + ], +) +def test_explicit_invalid_basis_is_rejected(invalid): + with pytest.raises(ValueError): + resolve({**COMPLETE_BASIS, **invalid}) + + +def test_survivor_benchmark_must_match_supplied_linked_computation_year(): + with pytest.raises(ValueError, match="computation years disagree"): + resolve(eligibility_year=2024) + + +@pytest.mark.parametrize("worker_index", [None, -1, 2]) +def test_invalid_worker_link_is_rejected(worker_index): + with pytest.raises(ValueError, match="existing person"): + resolve(worker_index=worker_index) diff --git a/policyengine_us/tools/social_security/population.py b/policyengine_us/tools/social_security/population.py new file mode 100644 index 00000000000..3a5f7bdfdf7 --- /dev/null +++ b/policyengine_us/tools/social_security/population.py @@ -0,0 +1,341 @@ +"""Reported-benefit PIA inference for independent annual cross-sections. + +Only unreformed reported inputs enter this adapter. Calculated sources already +read the shared formula path themselves and bypass the population ratio. Annual +survey amounts are treated as full-year gross benefits; their PIA intervals are +conditional diagnostics, not identified monthly-payment histories. +""" + +from datetime import date +from decimal import Decimal +from fractions import Fraction + +import numpy as np + +from policyengine_us.tools.social_security.cola import ( + cola_factor, + read_parameter_events, +) +from policyengine_us.tools.social_security.household import ( + attainment_month, + claiming_factor, + full_retirement_age_months, + month_string, +) +from policyengine_us.tools.social_security.parameters import ( + bend_points_at, + formula_path_at, +) +from policyengine_us.tools.social_security.pia import has_active_reform +from policyengine_us.tools.social_security.pia_ratio import compute_pia_ratio +from policyengine_us.tools.social_security.records import resolve_worker_basis +from policyengine_us.tools.social_security.population_survivor import ( + survivor_formula_path, +) +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, +) + + +SCENARIOS = { + "supplied_claim_age", + "entitlement_dates", + "age_62", + "age_65", + "fra", + "age_70", +} + + +def _choice(value): + # Core supports lists but not scalar strings. Retain readable named + # scenarios without adding another policy-selection mechanism. + if isinstance(value, (list, tuple, np.ndarray)) and len(value) == 1: + value = value[0] + if not isinstance(value, str): + raise ValueError("A PIA-ratio scenario must contain exactly one name") + return value + + +def _disability_proxy_year(parameters, year, scenario): + if scenario == "current_year": + return year + if scenario != "pre_reform": + raise ValueError("Choose pre_reform or current_year for disability eligibility") + # This is an exposure sensitivity, not an observed onset distribution. + # Use the cohort immediately before the first enabled formula change. + starts = [year] + node = parameters.gov.contrib.ssa.pia + for entry in node.in_effect.values_list: + if entry.value and int(entry.instant_str[:4]) <= year: + starts.append(int(entry.instant_str[:4])) + for entry in parameters.gov.contrib.congress.s770.section_2.in_effect.values_list: + if entry.value and int(entry.instant_str[:4]) <= year: + starts.append(max(2026, int(entry.instant_str[:4]))) + return max(1983, min(starts) - 1) + + +def _record_basis(person, period, parameters, component): + """Resolve supplied record information before any labelled age proxy.""" + prefix = f"social_security_{component}" + year = period.start.year + amounts = person(f"{prefix}_reported", period) + reported = ( + person(f"{prefix}_source", period) == SocialSecurityBenefitSource.reported + ) + if np.any(reported & (~np.isfinite(amounts) | (amounts < 0))): + raise ValueError( + "Reported Social Security benefits must be finite and nonnegative" + ) + positive = reported & (amounts > 0) + n = len(amounts) + age = person("age", period) + birth_years = person("birth_year", period) + years, linked, worker_indices = resolve_worker_basis( + person, period, component, positive + ) + supplied = years > 0 + entitlement = person(f"{prefix}_entitlement_year", period) + proxy = np.zeros(n, dtype=bool) + missing_claim_age = np.zeros(n, dtype=bool) + unmodeled = np.zeros(n, dtype=bool) + own = np.zeros(n, dtype=bool) + adjustments = [Fraction(1)] * n + shares = [Fraction(1)] * n + policy = parameters.gov.contrib.ssa.benefit_levers.pia_ratio(period.start) + + if component == "retirement": + # An entitlement date is a claim date, not age-62 eligibility. In its + # absence the scoring year's age rebuilds the cohort independently. + # A link without an explicit computation basis is not permission to + # infer an unrelated worker's record from the recipient's age. + unmodeled |= linked & ~supplied + own = positive & (age >= 62) & (worker_indices == np.arange(n)) + mixed = np.zeros(n, dtype=bool) + for other in ("survivors", "dependents", "disability"): + mixed |= person(f"social_security_{other}_reported", period) > 0 + own &= ~mixed + # Mixed recipient components do not establish a separable own-worker + # PIA in the presence of dual entitlement or partial-year conversions. + unmodeled |= positive & mixed + scenario = _choice(policy.claiming_scenario) + if scenario not in SCENARIOS: + raise ValueError("Unsupported PIA-ratio claiming_scenario") + months = person("social_security_retirement_claim_age_months", period) + birth_months = person("social_security_birth_month", period) + birth_days = person("social_security_birth_day", period) + birth_dates = person("ss_calculator_birth_date", period) + entitlement_months = person(f"{prefix}_entitlement_month", period) + births = [None] * n + exact_births = np.zeros(n, dtype=bool) + for i in np.flatnonzero(positive): + j = worker_indices[i] + birth = date(int(birth_years[j]), 1, 2) + exact_birth = False + if str(birth_dates[j]): + birth = date.fromisoformat(str(birth_dates[j])) + exact_birth = True + elif birth_months[j] and birth_days[j]: + try: + birth = date( + int(birth_years[j]), int(birth_months[j]), int(birth_days[j]) + ) + except ValueError as error: + raise ValueError( + "Invalid Social Security birth month/day" + ) from error + exact_birth = True + births[i], exact_births[i] = birth, exact_birth + if not supplied[i] and not linked[i]: + if exact_birth: + years[i] = attainment_month(birth, 744) // 12 + elif entitlement[i] > 0: + # A known claim age plus entitlement month identifies the + # worker's age-62 year without treating claiming as eligibility. + assumed_months = { + "supplied_claim_age": int(months[i]), + "age_62": 744, + "age_65": 780, + "fra": full_retirement_age_months(birth), + "age_70": 840, + "entitlement_dates": 0, + }[scenario] + if 744 <= assumed_months <= 840: + if not 1 <= entitlement_months[i] <= 12: + raise ValueError( + "Social Security entitlement month must lie in [1, 12]" + ) + years[i] = ( + int(entitlement[i]) * 12 + + int(entitlement_months[i]) + - 1 + - assumed_months + + 744 + ) // 12 + if not years[i]: + years[i] = year - age[i] + 62 + proxy[i] = True + for i in np.flatnonzero(positive & (years >= 1983)): + birth = births[i] + exact_birth = exact_births[i] + if not exact_birth: + proxy[i] = True + if scenario == "entitlement_dates": + if not exact_birth or entitlement[i] <= 0: + missing_claim_age[i] = True + continue + if not 1 <= entitlement_months[i] <= 12: + raise ValueError( + "Social Security entitlement month must lie in [1, 12]" + ) + claim = int(entitlement[i]) * 12 + int(entitlement_months[i]) - 1 + claim_months = claim - attainment_month(birth, 0) + else: + claim_months = { + "supplied_claim_age": int(months[i]), + "age_62": 744, + "age_65": 780, + "fra": full_retirement_age_months(birth), + "age_70": 840, + }[scenario] + claim = attainment_month(birth, claim_months) + if not 744 <= claim_months <= 840: + missing_claim_age[i] = True + continue + # The annual in-payment approximation uses fully credited DRCs. + # Partial entitlement years/deferred same-year credits require the + # household monthly ledger rather than annual survey inversion. + adjustments[i] = claiming_factor( + birth, month_string(claim), month_string((claim // 12 + 1) * 12) + ) + elif component == "disability": + years = np.where(years > 0, years, entitlement) + missing = years <= 0 + years = np.where( + missing, + _disability_proxy_year( + parameters, year, _choice(policy.disability_eligibility_proxy) + ), + years, + ) + proxy |= missing + else: + # A worker link/year and an identified composite a*s are both needed. + # No auxiliary inversion ever substitutes the recipient's age. + ratio = person(f"{prefix}_worker_benefit_ratio", period) + identified = person(f"{prefix}_worker_ratio_identified", period) + if np.any(positive & identified & (~np.isfinite(ratio) | (ratio <= 0))): + raise ValueError( + "An identified worker benefit ratio must be positive and finite" + ) + unmodeled |= ~(linked & supplied & identified) + for i in np.flatnonzero(positive & ~unmodeled): + shares[i] = Fraction(str(ratio[i])) + + if np.any(positive & supplied & (years > year)): + raise ValueError("Social Security eligibility cannot follow the scoring year") + unmodeled |= positive & (years > year) + pre1983 = positive & (years > 0) & (years < 1983) + unmodeled = positive & (unmodeled | missing_claim_age | pre1983) + survivor_paths = {} + if component == "survivors": + for i in np.flatnonzero(positive & ~unmodeled): + path, mode = formula_path_at(parameters, int(years[i]), year) + path, unsupported = survivor_formula_path( + person, period, parameters, i, path, int(years[i]), worker_indices[i] + ) + unmodeled[i] |= unsupported + survivor_paths[i] = (path, mode) + return dict( + amounts=amounts, + positive=positive, + eligibility_year=np.where(positive, years, 0), + claiming_adjustment=adjustments, + shares=shares, + own_record_retired_worker=own, + missing_claim_age=positive & missing_claim_age, + pre_1983=pre1983, + unmodeled=unmodeled, + proxy=positive & proxy, + measurement_uncertainty=positive & ~unmodeled, + survivor_paths=survivor_paths, + ) + + +def reported_pia_ratio_output(person, period, parameters, component, field): + """Named Person diagnostics and factor using one unreformed source basis.""" + if field == "factor": + year = period.start.year + custom = any( + entry.value and int(entry.instant_str[:4]) <= year + for entry in parameters.gov.contrib.ssa.pia.in_effect.values_list + ) + s770 = parameters.gov.contrib.congress.s770.section_2(period.start).in_effect + if not custom and not s770: + # Neutral current law must not require newly supplied record fields. + # This also preserves legacy amounts for records outside C's scope. + amount = person(f"social_security_{component}_reported", period) + reported = ( + person(f"social_security_{component}_source", period) + == SocialSecurityBenefitSource.reported + ) + if np.any(reported & (~np.isfinite(amount) | (amount < 0))): + raise ValueError( + "Reported Social Security benefits must be finite and nonnegative" + ) + return np.ones(len(amount)) + basis = _record_basis(person, period, parameters, component) + if field in basis and field != "claiming_adjustment": + return basis[field] + if field == "claiming_adjustment": + return np.array([float(value) for value in basis[field]]) + if field not in ("factor", "aime", "aime_lower", "aime_upper"): + raise ValueError("Unknown reported PIA ratio output") + result = ( + np.ones(len(basis["amounts"])) + if field == "factor" + else np.zeros(len(basis["amounts"])) + ) + supported = basis["positive"] & ~basis["unmodeled"] + if not np.any(supported): + return result + year = period.start.year + paths = {} + for eligibility in np.unique(basis["eligibility_year"][supported]): + eligibility = int(eligibility) + path, mode = formula_path_at(parameters, eligibility, year) + paths[eligibility] = (path, mode) + if field == "factor" and not any( + has_active_reform(path, eligibility, year, mode) + for eligibility, (path, mode) in paths.items() + ): + return result + events = read_parameter_events(parameters, year) + cola_by_year = {} + bends_by_year = {} + bound = parameters.gov.contrib.ssa.benefit_levers.pia_ratio( + period.start + ).annual_rounding_bound + for i in np.flatnonzero(supported): + eligibility = int(basis["eligibility_year"][i]) + if eligibility not in cola_by_year: + cola_by_year[eligibility] = cola_factor(eligibility, f"{year}-12", events) + bends_by_year[eligibility] = bend_points_at(parameters, eligibility) + path, mode = basis["survivor_paths"].get(i, paths[eligibility]) + # The shortest decimal representation of Core's float32 input restores + # ordinary supplied cents without expanding binary noise into a PIA dime. + amount = Decimal(str(basis["amounts"][i])) + inference = compute_pia_ratio( + amount, + eligibility, + path, + year, + mode, + bend_points=bends_by_year[eligibility], + claiming_adjustment=basis["claiming_adjustment"][i], + auxiliary_share=basis["shares"][i], + cola_factor=cola_by_year[eligibility], + annual_rounding_bound=bound, + ) + result[i] = float(getattr(inference, field)) + return result diff --git a/policyengine_us/tools/social_security/population_survivor.py b/policyengine_us/tools/social_security/population_survivor.py new file mode 100644 index 00000000000..508d67b0bf9 --- /dev/null +++ b/policyengine_us/tools/social_security/population_survivor.py @@ -0,0 +1,138 @@ +"""S.770 scope for an identified reported survivor worker record. + +This adapter reuses the supplied-record calculator's SurvivorBasis. Calculator +basis fields belong to the survivor recipient, just as on the household route; +a worker link does not authorize filling missing fields from recipient age or +combining partial bases from different people. Missing basis information leaves +the population amount unmodeled under S.770. Explicit contradictions are errors. +""" + +from dataclasses import replace +from datetime import date + +from policyengine_us.tools.social_security.household import ( + SurvivorBasis, + attainment_month, +) + + +def survivor_formula_path( + person, period, parameters, index, path, eligibility_year, worker_index +): + """Return (path, unsupported) with the shared survivor bend-point reading. + + Only an active S.770 path needs these additional death/computation fields. + Other custom formula reforms keep the caller's identified generic composite + worker-benefit ratio route. For S.770, the same ss_calculator_* basis inputs + used by 4a-1 must be supplied on the recipient record. Their benchmark must + equal the independently supplied linked worker eligibility/computation year. + + The pre-2026 death-before-any-eligibility switch changes only the first + bend-point increase. S.770's 95% first factor remains active either way. + A worker with authoritative prior-disability eligibility receives both + increases, as on the household route, regardless of this switch. + """ + if not path.s770: + return path, False + if worker_index is None or not 0 <= int(worker_index) < len(person.ids): + raise ValueError("Survivor worker link must identify an existing person") + + def value(name): + scalar = person("ss_calculator_" + name, period)[index] + return scalar.item() if hasattr(scalar, "item") else scalar + + def optional_date(name): + raw = str(value(name)) + return date.fromisoformat(raw) if raw else None + + worker_birth_year = int(value("worker_birth_year")) + death_year = int(value("death_year")) + authoritative_year = int(value("authoritative_year")) + prior_disability = int(value("prior_disability_eligibility")) + same_year = int(value("same_year_death_before_62")) + worker_birth_date = optional_date("worker_birth_date") + death_date = optional_date("death_date") + provenance = str(value("basis_provenance")) + ordinary_non_windex = bool(value("ordinary_non_windex")) + not_january_1 = bool(value("year_only_not_january_1")) + + if worker_birth_year < 0 or death_year < 0 or authoritative_year < 0: + raise ValueError( + "Survivor computation years must be positive or zero if missing" + ) + if prior_disability not in (-1, 0, 1): + raise ValueError("Survivor prior-disability assertion must be -1, 0 or 1") + if same_year not in (-1, 0, 1): + raise ValueError("Survivor same-year death assertion must be -1, 0 or 1") + if ( + worker_birth_date + and worker_birth_year + and worker_birth_date.year != worker_birth_year + ): + raise ValueError("Worker birth date and birth year disagree") + if death_date and death_year and death_date.year != death_year: + raise ValueError("Death date and death year disagree") + if death_date and worker_birth_date and death_date < worker_birth_date: + raise ValueError("Death cannot precede birth") + if worker_birth_year and death_year and death_year < worker_birth_year: + raise ValueError("Death cannot precede birth") + if death_year > period.start.year: + raise ValueError( + "A reported survivor's worker death cannot follow the scoring year" + ) + + # Default sentinels express unavailable information, not an assertion of a + # complete supported basis. Do not reject ordinary population artifacts just + # because they lack the household route's detailed scope declarations. + if ( + not worker_birth_year + or not death_year + or prior_disability == -1 + or not provenance.strip() + or not ordinary_non_windex + or (prior_disability == 1 and not authoritative_year) + or (not worker_birth_date and not not_january_1 and not authoritative_year) + ): + return path, True + + # A same-year death needs dates or the explicit assertion required by the + # shared service. With no such observation, this is missing identification. + year62 = ( + attainment_month(worker_birth_date, 744) // 12 + if worker_birth_date + else worker_birth_year + 62 + ) + if ( + death_year == year62 + and not (worker_birth_date and death_date) + and same_year == -1 + ): + return path, True + + basis = SurvivorBasis( + worker_birth_year=worker_birth_year, + death_year=death_year, + provenance=provenance, + worker_birth_date=worker_birth_date, + death_date=death_date, + authoritative_year=authoritative_year or None, + prior_disability_eligibility=bool(prior_disability), + ordinary_non_windex=ordinary_non_windex, + year_only_not_january_1=not_january_1, + same_year_death_before_62=None if same_year == -1 else bool(same_year), + ) + benchmark, before62 = basis.benchmark() + if benchmark != int(eligibility_year): + raise ValueError("Survivor basis and linked worker computation years disagree") + + restricted = ( + before62 and death_year < 2026 and not basis.prior_disability_eligibility + ) + if restricted: + increase = bool( + parameters.gov.contrib.congress.s770.section_2( + period.start + ).recompute_pre2026_death_before_eligibility + ) + return replace(path, s770_increase_first_bend_point=increase), False + return path, False diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_birth_day.py b/policyengine_us/variables/gov/ssa/ss/social_security_birth_day.py new file mode 100644 index 00000000000..8b70e458234 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_birth_day.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_birth_day(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security birth day" + documentation = "Supplied day of birth for the entitlement_dates population claiming scenario. Zero means unavailable. No formula or uprating. Exact dates preserve SSA age attainment on the day before a birthday, including January 1 births." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_birth_month.py b/policyengine_us/variables/gov/ssa/ss/social_security_birth_month.py new file mode 100644 index 00000000000..c5a9ac05e1a --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_birth_month.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_birth_month(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security birth month" + documentation = "Supplied birth month, 1 through 12, for the entitlement_dates population claiming scenario. Zero means unavailable. No formula or uprating. Used with the birth year and supplied birth day; an exact ss_calculator_birth_date can supply the complete date instead." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_inferred_aime.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_inferred_aime.py new file mode 100644 index 00000000000..277ff53e080 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_inferred_aime.py @@ -0,0 +1,16 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_dependents_inferred_aime(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = "Social Security dependents inferred eligibility-year monthly aime" + documentation = "Baseline AIME inverse at the observed unreformed amount, using baseline claiming, identified worker ratio and baseline COLAs. This is the lower endpoint for an exact dime-rounded eligibility PIA. A configured annual measurement bound widens the interval without changing this point estimate. Zero denotes no inference; inspect the diagnostic flags." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "dependents", "aime" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_inferred_aime_lower_bound.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_inferred_aime_lower_bound.py new file mode 100644 index 00000000000..517693f6588 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_inferred_aime_lower_bound.py @@ -0,0 +1,18 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_dependents_inferred_aime_lower_bound(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = ( + "Social Security dependents inferred eligibility-year monthly aime lower bound" + ) + documentation = "Inclusive lower endpoint of the conditional baseline-inversion interval in monthly AIME dollars. This is a PIA-rounding interval, not a confidence interval for survey measurement error." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "dependents", "aime_lower" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_inferred_aime_upper_bound.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_inferred_aime_upper_bound.py new file mode 100644 index 00000000000..975cde010e3 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_inferred_aime_upper_bound.py @@ -0,0 +1,18 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_dependents_inferred_aime_upper_bound(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = ( + "Social Security dependents inferred eligibility-year monthly aime upper bound" + ) + documentation = "Exclusive upper endpoint of the conditional baseline-inversion interval in monthly AIME dollars. Any configured annual rounding bound is propagated backward; additional survey measurement error is not quantified." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "dependents", "aime_upper" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_claiming_adjustment.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_claiming_adjustment.py new file mode 100644 index 00000000000..6180e5efe62 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_claiming_adjustment.py @@ -0,0 +1,16 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_dependents_pia_ratio_claiming_adjustment(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = "/1" + label = "Social Security dependents pia-ratio claiming adjustment" + documentation = "Worker claiming adjustment used in the baseline inversion under the explicitly selected scenario. Disability uses one. An auxiliary composite worker-benefit ratio already includes its claiming adjustment and is not divided twice." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "dependents", "claiming_adjustment" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_eligibility_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_eligibility_year.py new file mode 100644 index 00000000000..da8dc1e8ff6 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_eligibility_year.py @@ -0,0 +1,16 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_dependents_pia_ratio_eligibility_year(Variable): + value_type = int + entity = Person + definition_period = YEAR + unit = "year" + label = "Social Security dependents pia-ratio eligibility year" + documentation = "Underlying worker eligibility or computation year used by the population PIA-ratio calculation. Supplied record and worker links precede a labelled proxy; zero denotes an unavailable basis. Rebuilt independently in every scoring year." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "dependents", "eligibility_year" + ) 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 034b50023c9..fa929073ed8 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 @@ -1,4 +1,5 @@ from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output class social_security_dependents_pia_ratio_factor(Variable): @@ -8,8 +9,12 @@ class social_security_dependents_pia_ratio_factor(Variable): unit = "/1" label = "Social Security dependents PIA ratio factor" documentation = ( - "Neutral extension point for a subsequent reported-benefit PIA reform." + "Eligibility-year reformed-to-baseline PIA ratio for the reported source. " + "Calculated sources bypass this factor. Unsupported or unmodeled records " + "remain at one and expose separate diagnostic indicators." ) def formula(person, period, parameters): - return 1 + return reported_pia_ratio_output( + person, period, parameters, "dependents", "factor" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_measurement_uncertainty.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_measurement_uncertainty.py new file mode 100644 index 00000000000..f56d024044a --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_measurement_uncertainty.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_dependents_pia_ratio_measurement_uncertainty(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security dependents pia-ratio measurement uncertainty indicator" + documentation = "The inference uses a reported annual benefit whose payment or survey measurement basis is not an exact eligibility-level PIA. Its conditional rounding interval does not quantify all measurement uncertainty." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "dependents", "measurement_uncertainty" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_measurement_uncertainty_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_measurement_uncertainty_count.py new file mode 100644 index 00000000000..a4fdedb5ecf --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_measurement_uncertainty_count.py @@ -0,0 +1,22 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_pia_ratio_measurement_uncertainty_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security dependents PIA-ratio measurement uncertainty count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members( + "social_security_dependents_pia_ratio_measurement_uncertainty", period + ) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_missing_claim_age.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_missing_claim_age.py new file mode 100644 index 00000000000..b4a9dec1ce8 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_missing_claim_age.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_dependents_pia_ratio_missing_claim_age(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security dependents pia-ratio missing or unusable claiming data" + documentation = "Positive reported amount lacks usable data required by the selected claiming scenario. Under supplied_claim_age, missing, before-age-62 or after-age-70 values leave the factor at one without switching scenarios." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "dependents", "missing_claim_age" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_missing_claim_age_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_missing_claim_age_count.py new file mode 100644 index 00000000000..a606d474d9d --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_missing_claim_age_count.py @@ -0,0 +1,22 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_pia_ratio_missing_claim_age_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security dependents PIA-ratio missing claim age count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members( + "social_security_dependents_pia_ratio_missing_claim_age", period + ) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_own_record_retired_worker.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_own_record_retired_worker.py new file mode 100644 index 00000000000..4d48333822e --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_own_record_retired_worker.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_dependents_pia_ratio_own_record_retired_worker(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security dependents pia-ratio own-record retired-worker indicator" + documentation = "Whether this component is identified as an own-record retired-worker amount for the population PIA-ratio approximation. This indicator does not establish insured status or validate survey benefit classification." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "dependents", "own_record_retired_worker" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_pre_1983.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_pre_1983.py new file mode 100644 index 00000000000..db62d691ec5 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_pre_1983.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_dependents_pia_ratio_pre_1983(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security dependents pia-ratio unsupported pre-1983 benchmark" + documentation = "Positive reported amount has a benchmark year before 1983. The population PIA-ratio factor remains one because the supported inverse assumes initial PIA dime truncation." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "dependents", "pre_1983" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_pre_1983_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_pre_1983_count.py new file mode 100644 index 00000000000..752ae974041 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_pre_1983_count.py @@ -0,0 +1,20 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_pia_ratio_pre_1983_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security dependents PIA-ratio pre 1983 count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members("social_security_dependents_pia_ratio_pre_1983", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_proxy.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_proxy.py new file mode 100644 index 00000000000..859df7dce22 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_proxy.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_dependents_pia_ratio_proxy(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security dependents pia-ratio proxy assumption indicator" + documentation = "The population PIA-ratio diagnostics use an eligibility, claiming, or birth-date approximation rather than a fully supplied record basis. This flag distinguishes sensitivity assumptions from observed history." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "dependents", "proxy" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_proxy_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_proxy_count.py new file mode 100644 index 00000000000..dd3c9e297c9 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_proxy_count.py @@ -0,0 +1,20 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_pia_ratio_proxy_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security dependents PIA-ratio proxy count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members("social_security_dependents_pia_ratio_proxy", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_unmodeled.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_unmodeled.py new file mode 100644 index 00000000000..98a3fb42047 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_unmodeled.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_dependents_pia_ratio_unmodeled(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security dependents pia-ratio unmodeled amount indicator" + documentation = "Positive reported amount lacks an identified worker basis, composite ratio, or another input required for a supported inversion. The factor stays one; the amount is unmodeled for this lever." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "dependents", "unmodeled" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_unmodeled_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_unmodeled_count.py new file mode 100644 index 00000000000..972d70a04d0 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_unmodeled_count.py @@ -0,0 +1,20 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_pia_ratio_unmodeled_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security dependents PIA-ratio unmodeled count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members("social_security_dependents_pia_ratio_unmodeled", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_benefit_ratio.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_benefit_ratio.py new file mode 100644 index 00000000000..3179d5bd3fb --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_benefit_ratio.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_worker_benefit_ratio(Variable): + value_type = float + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security dependents composite worker-benefit ratio" + documentation = "Supplied positive composite ratio of this recipient benefit to the underlying worker PIA, including the worker claiming adjustment and all applicable recipient share reductions. The baseline inversion divides by this composite once. Zero means unavailable. Requires an identified worker basis and the separate worker_ratio_identified assertion; no uprating." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_ratio_identified.py b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_ratio_identified.py new file mode 100644 index 00000000000..aceeb2f3dae --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_worker_ratio_identified.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_dependents_worker_ratio_identified(Variable): + value_type = bool + default_value = False + entity = Person + definition_period = YEAR + label = "Social Security dependents worker ratio identified" + documentation = "Assert that the supplied composite worker-benefit ratio identifies a single underlying worker PIA, including all reductions and credits, with no unaccounted dual entitlement, binding family cap, or earnings-test withholding. False leaves this auxiliary amount unmodeled for the population PIA-ratio lever." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_inferred_aime.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_inferred_aime.py new file mode 100644 index 00000000000..89b039e9553 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_inferred_aime.py @@ -0,0 +1,16 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_disability_inferred_aime(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = "Social Security disability inferred eligibility-year monthly aime" + documentation = "Baseline AIME inverse at the observed unreformed amount, using baseline claiming, identified worker ratio and baseline COLAs. This is the lower endpoint for an exact dime-rounded eligibility PIA. A configured annual measurement bound widens the interval without changing this point estimate. Zero denotes no inference; inspect the diagnostic flags." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "disability", "aime" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_inferred_aime_lower_bound.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_inferred_aime_lower_bound.py new file mode 100644 index 00000000000..950ffa6f2cf --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_inferred_aime_lower_bound.py @@ -0,0 +1,18 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_disability_inferred_aime_lower_bound(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = ( + "Social Security disability inferred eligibility-year monthly aime lower bound" + ) + documentation = "Inclusive lower endpoint of the conditional baseline-inversion interval in monthly AIME dollars. This is a PIA-rounding interval, not a confidence interval for survey measurement error." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "disability", "aime_lower" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_inferred_aime_upper_bound.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_inferred_aime_upper_bound.py new file mode 100644 index 00000000000..09074047fb3 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_inferred_aime_upper_bound.py @@ -0,0 +1,18 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_disability_inferred_aime_upper_bound(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = ( + "Social Security disability inferred eligibility-year monthly aime upper bound" + ) + documentation = "Exclusive upper endpoint of the conditional baseline-inversion interval in monthly AIME dollars. Any configured annual rounding bound is propagated backward; additional survey measurement error is not quantified." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "disability", "aime_upper" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_claiming_adjustment.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_claiming_adjustment.py new file mode 100644 index 00000000000..ad13035523e --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_claiming_adjustment.py @@ -0,0 +1,16 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_disability_pia_ratio_claiming_adjustment(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = "/1" + label = "Social Security disability pia-ratio claiming adjustment" + documentation = "Worker claiming adjustment used in the baseline inversion under the explicitly selected scenario. Disability uses one. An auxiliary composite worker-benefit ratio already includes its claiming adjustment and is not divided twice." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "disability", "claiming_adjustment" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_eligibility_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_eligibility_year.py new file mode 100644 index 00000000000..6fb406f9a8a --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_eligibility_year.py @@ -0,0 +1,16 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_disability_pia_ratio_eligibility_year(Variable): + value_type = int + entity = Person + definition_period = YEAR + unit = "year" + label = "Social Security disability pia-ratio eligibility year" + documentation = "Underlying worker eligibility or computation year used by the population PIA-ratio calculation. Supplied record and worker links precede a labelled proxy; zero denotes an unavailable basis. Rebuilt independently in every scoring year." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "disability", "eligibility_year" + ) 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 e807b247634..1aa38473f6c 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 @@ -1,4 +1,5 @@ from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output class social_security_disability_pia_ratio_factor(Variable): @@ -8,8 +9,12 @@ class social_security_disability_pia_ratio_factor(Variable): unit = "/1" label = "Social Security disability PIA ratio factor" documentation = ( - "Neutral extension point for a subsequent reported-benefit PIA reform." + "Eligibility-year reformed-to-baseline PIA ratio for the reported source. " + "Calculated sources bypass this factor. Unsupported or unmodeled records " + "remain at one and expose separate diagnostic indicators." ) def formula(person, period, parameters): - return 1 + return reported_pia_ratio_output( + person, period, parameters, "disability", "factor" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_measurement_uncertainty.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_measurement_uncertainty.py new file mode 100644 index 00000000000..8f74d309533 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_measurement_uncertainty.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_disability_pia_ratio_measurement_uncertainty(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security disability pia-ratio measurement uncertainty indicator" + documentation = "The inference uses a reported annual benefit whose payment or survey measurement basis is not an exact eligibility-level PIA. Its conditional rounding interval does not quantify all measurement uncertainty." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "disability", "measurement_uncertainty" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_measurement_uncertainty_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_measurement_uncertainty_count.py new file mode 100644 index 00000000000..103de0176c6 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_measurement_uncertainty_count.py @@ -0,0 +1,22 @@ +from policyengine_us.model_api import * + + +class social_security_disability_pia_ratio_measurement_uncertainty_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security disability PIA-ratio measurement uncertainty count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members( + "social_security_disability_pia_ratio_measurement_uncertainty", period + ) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_missing_claim_age.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_missing_claim_age.py new file mode 100644 index 00000000000..e0030573a7e --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_missing_claim_age.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_disability_pia_ratio_missing_claim_age(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security disability pia-ratio missing or unusable claiming data" + documentation = "Positive reported amount lacks usable data required by the selected claiming scenario. Under supplied_claim_age, missing, before-age-62 or after-age-70 values leave the factor at one without switching scenarios." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "disability", "missing_claim_age" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_missing_claim_age_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_missing_claim_age_count.py new file mode 100644 index 00000000000..36cc4c11331 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_missing_claim_age_count.py @@ -0,0 +1,22 @@ +from policyengine_us.model_api import * + + +class social_security_disability_pia_ratio_missing_claim_age_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security disability PIA-ratio missing claim age count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members( + "social_security_disability_pia_ratio_missing_claim_age", period + ) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_own_record_retired_worker.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_own_record_retired_worker.py new file mode 100644 index 00000000000..71ca9ff4095 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_own_record_retired_worker.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_disability_pia_ratio_own_record_retired_worker(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security disability pia-ratio own-record retired-worker indicator" + documentation = "Whether this component is identified as an own-record retired-worker amount for the population PIA-ratio approximation. This indicator does not establish insured status or validate survey benefit classification." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "disability", "own_record_retired_worker" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_pre_1983.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_pre_1983.py new file mode 100644 index 00000000000..9152a1fb02e --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_pre_1983.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_disability_pia_ratio_pre_1983(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security disability pia-ratio unsupported pre-1983 benchmark" + documentation = "Positive reported amount has a benchmark year before 1983. The population PIA-ratio factor remains one because the supported inverse assumes initial PIA dime truncation." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "disability", "pre_1983" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_pre_1983_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_pre_1983_count.py new file mode 100644 index 00000000000..03be6d99f71 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_pre_1983_count.py @@ -0,0 +1,20 @@ +from policyengine_us.model_api import * + + +class social_security_disability_pia_ratio_pre_1983_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security disability PIA-ratio pre 1983 count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members("social_security_disability_pia_ratio_pre_1983", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_proxy.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_proxy.py new file mode 100644 index 00000000000..fe78b9b41a8 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_proxy.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_disability_pia_ratio_proxy(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security disability pia-ratio proxy assumption indicator" + documentation = "The population PIA-ratio diagnostics use an eligibility, claiming, or birth-date approximation rather than a fully supplied record basis. This flag distinguishes sensitivity assumptions from observed history." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "disability", "proxy" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_proxy_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_proxy_count.py new file mode 100644 index 00000000000..ae60d24a22f --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_proxy_count.py @@ -0,0 +1,20 @@ +from policyengine_us.model_api import * + + +class social_security_disability_pia_ratio_proxy_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security disability PIA-ratio proxy count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members("social_security_disability_pia_ratio_proxy", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_unmodeled.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_unmodeled.py new file mode 100644 index 00000000000..34e72d5b920 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_unmodeled.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_disability_pia_ratio_unmodeled(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security disability pia-ratio unmodeled amount indicator" + documentation = "Positive reported amount lacks an identified worker basis, composite ratio, or another input required for a supported inversion. The factor stays one; the amount is unmodeled for this lever." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "disability", "unmodeled" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_unmodeled_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_unmodeled_count.py new file mode 100644 index 00000000000..41730bebe0d --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_unmodeled_count.py @@ -0,0 +1,20 @@ +from policyengine_us.model_api import * + + +class social_security_disability_pia_ratio_unmodeled_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security disability PIA-ratio unmodeled count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members("social_security_disability_pia_ratio_unmodeled", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_claim_age_months.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_claim_age_months.py new file mode 100644 index 00000000000..feb7dd21d5c --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_claim_age_months.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_claim_age_months(Variable): + value_type = int + default_value = 0 + entity = Person + definition_period = YEAR + label = "Retirement age in whole months at first entitlement" + documentation = "Pure supplied retirement claim age in whole months: 62 years 3 months is 747. Zero means unavailable. No formula or uprating; projected dataset columns must supply the value separately for each scoring year. The selected supplied_claim_age scenario accepts ages 62 through 70." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_inferred_aime.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_inferred_aime.py new file mode 100644 index 00000000000..b97c8000a0f --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_inferred_aime.py @@ -0,0 +1,16 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_retirement_inferred_aime(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = "Social Security retirement inferred eligibility-year monthly aime" + documentation = "Baseline AIME inverse at the observed unreformed amount, using baseline claiming, identified worker ratio and baseline COLAs. This is the lower endpoint for an exact dime-rounded eligibility PIA. A configured annual measurement bound widens the interval without changing this point estimate. Zero denotes no inference; inspect the diagnostic flags." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "retirement", "aime" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_inferred_aime_lower_bound.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_inferred_aime_lower_bound.py new file mode 100644 index 00000000000..8024a5dc5f1 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_inferred_aime_lower_bound.py @@ -0,0 +1,18 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_retirement_inferred_aime_lower_bound(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = ( + "Social Security retirement inferred eligibility-year monthly aime lower bound" + ) + documentation = "Inclusive lower endpoint of the conditional baseline-inversion interval in monthly AIME dollars. This is a PIA-rounding interval, not a confidence interval for survey measurement error." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "retirement", "aime_lower" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_inferred_aime_upper_bound.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_inferred_aime_upper_bound.py new file mode 100644 index 00000000000..b3f488826fc --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_inferred_aime_upper_bound.py @@ -0,0 +1,18 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_retirement_inferred_aime_upper_bound(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = ( + "Social Security retirement inferred eligibility-year monthly aime upper bound" + ) + documentation = "Exclusive upper endpoint of the conditional baseline-inversion interval in monthly AIME dollars. Any configured annual rounding bound is propagated backward; additional survey measurement error is not quantified." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "retirement", "aime_upper" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_claiming_adjustment.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_claiming_adjustment.py new file mode 100644 index 00000000000..ac3c16279b0 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_claiming_adjustment.py @@ -0,0 +1,16 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_retirement_pia_ratio_claiming_adjustment(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = "/1" + label = "Social Security retirement pia-ratio claiming adjustment" + documentation = "Worker claiming adjustment used in the baseline inversion under the explicitly selected scenario. Disability uses one. An auxiliary composite worker-benefit ratio already includes its claiming adjustment and is not divided twice." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "retirement", "claiming_adjustment" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_eligibility_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_eligibility_year.py new file mode 100644 index 00000000000..9ba947d99df --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_eligibility_year.py @@ -0,0 +1,16 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_retirement_pia_ratio_eligibility_year(Variable): + value_type = int + entity = Person + definition_period = YEAR + unit = "year" + label = "Social Security retirement pia-ratio eligibility year" + documentation = "Underlying worker eligibility or computation year used by the population PIA-ratio calculation. Supplied record and worker links precede a labelled proxy; zero denotes an unavailable basis. Rebuilt independently in every scoring year." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "retirement", "eligibility_year" + ) 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 50375513765..02563922313 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 @@ -1,4 +1,5 @@ from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output class social_security_retirement_pia_ratio_factor(Variable): @@ -8,8 +9,12 @@ class social_security_retirement_pia_ratio_factor(Variable): unit = "/1" label = "Social Security retirement PIA ratio factor" documentation = ( - "Neutral extension point for a subsequent reported-benefit PIA reform." + "Eligibility-year reformed-to-baseline PIA ratio for the reported source. " + "Calculated sources bypass this factor. Unsupported or unmodeled records " + "remain at one and expose separate diagnostic indicators." ) def formula(person, period, parameters): - return 1 + return reported_pia_ratio_output( + person, period, parameters, "retirement", "factor" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_measurement_uncertainty.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_measurement_uncertainty.py new file mode 100644 index 00000000000..dcafd307c92 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_measurement_uncertainty.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_retirement_pia_ratio_measurement_uncertainty(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security retirement pia-ratio measurement uncertainty indicator" + documentation = "The inference uses a reported annual benefit whose payment or survey measurement basis is not an exact eligibility-level PIA. Its conditional rounding interval does not quantify all measurement uncertainty." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "retirement", "measurement_uncertainty" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_measurement_uncertainty_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_measurement_uncertainty_count.py new file mode 100644 index 00000000000..3d31811cfac --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_measurement_uncertainty_count.py @@ -0,0 +1,22 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_pia_ratio_measurement_uncertainty_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security retirement PIA-ratio measurement uncertainty count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members( + "social_security_retirement_pia_ratio_measurement_uncertainty", period + ) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_missing_claim_age.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_missing_claim_age.py new file mode 100644 index 00000000000..1e7b78c1e37 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_missing_claim_age.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_retirement_pia_ratio_missing_claim_age(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security retirement pia-ratio missing or unusable claiming data" + documentation = "Positive reported amount lacks usable data required by the selected claiming scenario. Under supplied_claim_age, missing, before-age-62 or after-age-70 values leave the factor at one without switching scenarios." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "retirement", "missing_claim_age" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_missing_claim_age_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_missing_claim_age_count.py new file mode 100644 index 00000000000..d044866684d --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_missing_claim_age_count.py @@ -0,0 +1,22 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_pia_ratio_missing_claim_age_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security retirement PIA-ratio missing claim age count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members( + "social_security_retirement_pia_ratio_missing_claim_age", period + ) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_own_record_retired_worker.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_own_record_retired_worker.py new file mode 100644 index 00000000000..9ffb037810d --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_own_record_retired_worker.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_retirement_pia_ratio_own_record_retired_worker(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security retirement pia-ratio own-record retired-worker indicator" + documentation = "Whether this component is identified as an own-record retired-worker amount for the population PIA-ratio approximation. This indicator does not establish insured status or validate survey benefit classification." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "retirement", "own_record_retired_worker" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_pre_1983.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_pre_1983.py new file mode 100644 index 00000000000..a20624f5477 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_pre_1983.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_retirement_pia_ratio_pre_1983(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security retirement pia-ratio unsupported pre-1983 benchmark" + documentation = "Positive reported amount has a benchmark year before 1983. The population PIA-ratio factor remains one because the supported inverse assumes initial PIA dime truncation." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "retirement", "pre_1983" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_pre_1983_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_pre_1983_count.py new file mode 100644 index 00000000000..bd55768e3f7 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_pre_1983_count.py @@ -0,0 +1,20 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_pia_ratio_pre_1983_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security retirement PIA-ratio pre 1983 count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members("social_security_retirement_pia_ratio_pre_1983", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_proxy.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_proxy.py new file mode 100644 index 00000000000..6753193f36a --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_proxy.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_retirement_pia_ratio_proxy(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security retirement pia-ratio proxy assumption indicator" + documentation = "The population PIA-ratio diagnostics use an eligibility, claiming, or birth-date approximation rather than a fully supplied record basis. This flag distinguishes sensitivity assumptions from observed history." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "retirement", "proxy" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_proxy_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_proxy_count.py new file mode 100644 index 00000000000..828d0fac5d9 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_proxy_count.py @@ -0,0 +1,20 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_pia_ratio_proxy_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security retirement PIA-ratio proxy count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members("social_security_retirement_pia_ratio_proxy", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_unmodeled.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_unmodeled.py new file mode 100644 index 00000000000..e1badf1a94a --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_unmodeled.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_retirement_pia_ratio_unmodeled(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security retirement pia-ratio unmodeled amount indicator" + documentation = "Positive reported amount lacks an identified worker basis, composite ratio, or another input required for a supported inversion. The factor stays one; the amount is unmodeled for this lever." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "retirement", "unmodeled" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_unmodeled_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_unmodeled_count.py new file mode 100644 index 00000000000..7f97b131663 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_unmodeled_count.py @@ -0,0 +1,20 @@ +from policyengine_us.model_api import * + + +class social_security_retirement_pia_ratio_unmodeled_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security retirement PIA-ratio unmodeled count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members("social_security_retirement_pia_ratio_unmodeled", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_inferred_aime.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_inferred_aime.py new file mode 100644 index 00000000000..0486bc7ea9d --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_inferred_aime.py @@ -0,0 +1,16 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_survivors_inferred_aime(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = "Social Security survivors inferred eligibility-year monthly aime" + documentation = "Baseline AIME inverse at the observed unreformed amount, using baseline claiming, identified worker ratio and baseline COLAs. This is the lower endpoint for an exact dime-rounded eligibility PIA. A configured annual measurement bound widens the interval without changing this point estimate. Zero denotes no inference; inspect the diagnostic flags." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "survivors", "aime" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_inferred_aime_lower_bound.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_inferred_aime_lower_bound.py new file mode 100644 index 00000000000..86b66878da9 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_inferred_aime_lower_bound.py @@ -0,0 +1,18 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_survivors_inferred_aime_lower_bound(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = ( + "Social Security survivors inferred eligibility-year monthly aime lower bound" + ) + documentation = "Inclusive lower endpoint of the conditional baseline-inversion interval in monthly AIME dollars. This is a PIA-rounding interval, not a confidence interval for survey measurement error." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "survivors", "aime_lower" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_inferred_aime_upper_bound.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_inferred_aime_upper_bound.py new file mode 100644 index 00000000000..c6e798a351f --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_inferred_aime_upper_bound.py @@ -0,0 +1,18 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_survivors_inferred_aime_upper_bound(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = USD + label = ( + "Social Security survivors inferred eligibility-year monthly aime upper bound" + ) + documentation = "Exclusive upper endpoint of the conditional baseline-inversion interval in monthly AIME dollars. Any configured annual rounding bound is propagated backward; additional survey measurement error is not quantified." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "survivors", "aime_upper" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_claiming_adjustment.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_claiming_adjustment.py new file mode 100644 index 00000000000..a350a1d790e --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_claiming_adjustment.py @@ -0,0 +1,16 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_survivors_pia_ratio_claiming_adjustment(Variable): + value_type = float + entity = Person + definition_period = YEAR + unit = "/1" + label = "Social Security survivors pia-ratio claiming adjustment" + documentation = "Worker claiming adjustment used in the baseline inversion under the explicitly selected scenario. Disability uses one. An auxiliary composite worker-benefit ratio already includes its claiming adjustment and is not divided twice." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "survivors", "claiming_adjustment" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_eligibility_year.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_eligibility_year.py new file mode 100644 index 00000000000..f987b4137d0 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_eligibility_year.py @@ -0,0 +1,16 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_survivors_pia_ratio_eligibility_year(Variable): + value_type = int + entity = Person + definition_period = YEAR + unit = "year" + label = "Social Security survivors pia-ratio eligibility year" + documentation = "Underlying worker eligibility or computation year used by the population PIA-ratio calculation. Supplied record and worker links precede a labelled proxy; zero denotes an unavailable basis. Rebuilt independently in every scoring year." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "survivors", "eligibility_year" + ) 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 40be5f4d21d..df21d801594 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 @@ -1,4 +1,5 @@ from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output class social_security_survivors_pia_ratio_factor(Variable): @@ -8,8 +9,12 @@ class social_security_survivors_pia_ratio_factor(Variable): unit = "/1" label = "Social Security survivors PIA ratio factor" documentation = ( - "Neutral extension point for a subsequent reported-benefit PIA reform." + "Eligibility-year reformed-to-baseline PIA ratio for the reported source. " + "Calculated sources bypass this factor. Unsupported or unmodeled records " + "remain at one and expose separate diagnostic indicators." ) def formula(person, period, parameters): - return 1 + return reported_pia_ratio_output( + person, period, parameters, "survivors", "factor" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_measurement_uncertainty.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_measurement_uncertainty.py new file mode 100644 index 00000000000..57f5d9485e7 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_measurement_uncertainty.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_survivors_pia_ratio_measurement_uncertainty(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security survivors pia-ratio measurement uncertainty indicator" + documentation = "The inference uses a reported annual benefit whose payment or survey measurement basis is not an exact eligibility-level PIA. Its conditional rounding interval does not quantify all measurement uncertainty." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "survivors", "measurement_uncertainty" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_measurement_uncertainty_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_measurement_uncertainty_count.py new file mode 100644 index 00000000000..78ac1378628 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_measurement_uncertainty_count.py @@ -0,0 +1,22 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_pia_ratio_measurement_uncertainty_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security survivors PIA-ratio measurement uncertainty count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members( + "social_security_survivors_pia_ratio_measurement_uncertainty", period + ) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_missing_claim_age.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_missing_claim_age.py new file mode 100644 index 00000000000..287f56960a4 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_missing_claim_age.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_survivors_pia_ratio_missing_claim_age(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security survivors pia-ratio missing or unusable claiming data" + documentation = "Positive reported amount lacks usable data required by the selected claiming scenario. Under supplied_claim_age, missing, before-age-62 or after-age-70 values leave the factor at one without switching scenarios." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "survivors", "missing_claim_age" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_missing_claim_age_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_missing_claim_age_count.py new file mode 100644 index 00000000000..05d8c3fb79c --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_missing_claim_age_count.py @@ -0,0 +1,22 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_pia_ratio_missing_claim_age_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security survivors PIA-ratio missing claim age count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members( + "social_security_survivors_pia_ratio_missing_claim_age", period + ) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_own_record_retired_worker.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_own_record_retired_worker.py new file mode 100644 index 00000000000..24e310580d5 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_own_record_retired_worker.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_survivors_pia_ratio_own_record_retired_worker(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security survivors pia-ratio own-record retired-worker indicator" + documentation = "Whether this component is identified as an own-record retired-worker amount for the population PIA-ratio approximation. This indicator does not establish insured status or validate survey benefit classification." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "survivors", "own_record_retired_worker" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_pre_1983.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_pre_1983.py new file mode 100644 index 00000000000..04293fc0629 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_pre_1983.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_survivors_pia_ratio_pre_1983(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security survivors pia-ratio unsupported pre-1983 benchmark" + documentation = "Positive reported amount has a benchmark year before 1983. The population PIA-ratio factor remains one because the supported inverse assumes initial PIA dime truncation." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "survivors", "pre_1983" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_pre_1983_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_pre_1983_count.py new file mode 100644 index 00000000000..9fe9c5864b0 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_pre_1983_count.py @@ -0,0 +1,20 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_pia_ratio_pre_1983_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security survivors PIA-ratio pre 1983 count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members("social_security_survivors_pia_ratio_pre_1983", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_proxy.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_proxy.py new file mode 100644 index 00000000000..ca2cabb8952 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_proxy.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_survivors_pia_ratio_proxy(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security survivors pia-ratio proxy assumption indicator" + documentation = "The population PIA-ratio diagnostics use an eligibility, claiming, or birth-date approximation rather than a fully supplied record basis. This flag distinguishes sensitivity assumptions from observed history." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "survivors", "proxy" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_proxy_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_proxy_count.py new file mode 100644 index 00000000000..9e27655f647 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_proxy_count.py @@ -0,0 +1,20 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_pia_ratio_proxy_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security survivors PIA-ratio proxy count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members("social_security_survivors_pia_ratio_proxy", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_unmodeled.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_unmodeled.py new file mode 100644 index 00000000000..bb389277848 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_unmodeled.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * +from policyengine_us.tools.social_security.population import reported_pia_ratio_output + + +class social_security_survivors_pia_ratio_unmodeled(Variable): + value_type = bool + entity = Person + definition_period = YEAR + label = "Social Security survivors pia-ratio unmodeled amount indicator" + documentation = "Positive reported amount lacks an identified worker basis, composite ratio, or another input required for a supported inversion. The factor stays one; the amount is unmodeled for this lever." + + def formula(person, period, parameters): + return reported_pia_ratio_output( + person, period, parameters, "survivors", "unmodeled" + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_unmodeled_count.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_unmodeled_count.py new file mode 100644 index 00000000000..fa6ac2eb767 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_unmodeled_count.py @@ -0,0 +1,20 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_pia_ratio_unmodeled_count(Variable): + value_type = float + entity = Household + definition_period = YEAR + unit = "person" + label = "Social Security survivors PIA-ratio unmodeled count" + documentation = ( + "Number of household members with the corresponding Person indicator. " + "Population aggregation uses household weights once. If supplied Person " + "weights differ within a household, use the Person indicator and its own " + "Person weights for the authoritative weighted count." + ) + + def formula(household, period, parameters): + return household.sum( + household.members("social_security_survivors_pia_ratio_unmodeled", period) + ) diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_benefit_ratio.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_benefit_ratio.py new file mode 100644 index 00000000000..6a682c8e8ff --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_benefit_ratio.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_worker_benefit_ratio(Variable): + value_type = float + default_value = 0 + entity = Person + definition_period = YEAR + label = "Social Security survivors composite worker-benefit ratio" + documentation = "Supplied positive composite ratio of this recipient benefit to the underlying worker PIA, including the worker claiming adjustment and all applicable recipient share reductions. The baseline inversion divides by this composite once. Zero means unavailable. Requires an identified worker basis and the separate worker_ratio_identified assertion; no uprating." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_ratio_identified.py b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_ratio_identified.py new file mode 100644 index 00000000000..ea8fef57b5d --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_worker_ratio_identified.py @@ -0,0 +1,10 @@ +from policyengine_us.model_api import * + + +class social_security_survivors_worker_ratio_identified(Variable): + value_type = bool + default_value = False + entity = Person + definition_period = YEAR + label = "Social Security survivors worker ratio identified" + documentation = "Assert that the supplied composite worker-benefit ratio identifies a single underlying worker PIA, including all reductions and credits, with no unaccounted dual entitlement, binding family cap, or earnings-test withholding. False leaves this auxiliary amount unmodeled for the population PIA-ratio lever." From 3b5fc99ff1e2cd480ed968d4d3141509651fcc85 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 27 Oct 2026 05:51:40 -0700 Subject: [PATCH 68/80] Bypass inference work for zero and calculated benefit records --- .../tools/social_security/population.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/policyengine_us/tools/social_security/population.py b/policyengine_us/tools/social_security/population.py index 3a5f7bdfdf7..0a4e74ae5e3 100644 --- a/policyengine_us/tools/social_security/population.py +++ b/policyengine_us/tools/social_security/population.py @@ -89,6 +89,21 @@ def _record_basis(person, period, parameters, component): ) positive = reported & (amounts > 0) n = len(amounts) + if not np.any(positive): + return dict( + amounts=amounts, + positive=positive, + eligibility_year=np.zeros(n, dtype=int), + claiming_adjustment=[Fraction(1)] * n, + shares=[Fraction(1)] * n, + own_record_retired_worker=np.zeros(n, dtype=bool), + missing_claim_age=np.zeros(n, dtype=bool), + pre_1983=np.zeros(n, dtype=bool), + unmodeled=np.zeros(n, dtype=bool), + proxy=np.zeros(n, dtype=bool), + measurement_uncertainty=np.zeros(n, dtype=bool), + survivor_paths={}, + ) age = person("age", period) birth_years = person("birth_year", period) years, linked, worker_indices = resolve_worker_basis( From 1297873dadf1973788da76cad12ec7126d599f5d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 27 Oct 2026 06:18:06 -0700 Subject: [PATCH 69/80] Keep population inference cohorts and rounded endpoints consistent --- docs/usage/social-security-pia-ratio.md | 48 +++++++-- .../tests/core/test_ss_pia_ratio_inputs.py | 50 +++++++++- ..._security_retirement_pia_ratio_factor.yaml | 55 ++++++++++- .../tools/social_security/population.py | 97 +++++++++++++------ .../gov/ssa/ss/social_security_birth_day.py | 2 +- .../gov/ssa/ss/social_security_birth_month.py | 2 +- ...etirement_pia_ratio_claiming_adjustment.py | 2 +- 7 files changed, 207 insertions(+), 49 deletions(-) diff --git a/docs/usage/social-security-pia-ratio.md b/docs/usage/social-security-pia-ratio.md index 66cf4b43114..d0f5bd7403e 100644 --- a/docs/usage/social-security-pia-ratio.md +++ b/docs/usage/social-security-pia-ratio.md @@ -55,15 +55,22 @@ All scenarios preserve cohort FRA, early months and delayed-credit rates. Under `supplied_claim_age`, a missing value, a value before 62, or a value after 70 leaves the affected factor at one and raises `_pia_ratio_missing_claim_age`. There is no silent change to another scenario. -The `entitlement_dates` scenario requires a known month and day of birth, -supplied as `social_security_birth_month` and `social_security_birth_day`, -or a complete `ss_calculator_birth_date`. Zero month/day means unavailable. -Exact dates preserve SSA attainment on the day before a birthday, including -January 1 births. +The `entitlement_dates` scenario uses the entitlement year/month and birth +year/month. Supply `social_security_birth_month`; the optional +`social_security_birth_day` refines SSA age attainment. With a known birth +month and no day, day 2 is the explicit within-month convention and +`_pia_ratio_proxy` is true. Zero birth month means unavailable. A complete +`ss_calculator_birth_date` supplies the exact date instead; exact dates +preserve attainment on the day before a birthday, including January 1 births. Where only a birth year is available for an age-month scenario, the cohort calculation uses the documented non-January-1 approximation and raises -`_pia_ratio_proxy`. A claiming scenario is substantive: in C2, replacing the +`_pia_ratio_proxy`. Exact birth information and an explicitly supplied +`birth_year` take precedence over a year derived from current age. When birth +information is unavailable but a worker eligibility year is supplied, E minus +62 can supply the approximate birth cohort for FRA and credit rates; this does +not change the supplied eligibility year. A claiming scenario is substantive: +in C2, replacing the correct age-62 adjustment of 0.70 with 1 infers AIME $2,978 instead of $5,000 and changes the annual gain from $1,165.248 to $1,664.64 before rounding, an overstatement of about 43%. @@ -77,9 +84,13 @@ entitlement/onset year supplies the component's applicable basis where appropriate. A link must identify an existing Person and a supplied computation basis; conflicting record years are unsupported. -For own retirement without a supplied basis, each scoring year's eligibility -proxy is `E = t - age_t + 62`. Rebuild it independently for every annual -cross-section. A row whose age stays 70 has E=2018 in 2026 and E=2027 in 2035; +For own retirement, a supplied worker eligibility year takes precedence. Without +one, exact birth information or an explicitly supplied `birth_year` identifies +the age-62 eligibility year before the current-age proxy; known entitlement +dates and the selected claim age can also identify that year. If none of these +supplies a basis, each scoring year's proxy is `E = t - age_t + 62`. Rebuild +it independently for every annual cross-section. A row whose age stays 70 and +has no supplied basis has E=2018 in 2026 and E=2027 in 2035; the model does not carry an inferred cohort or claiming history between those rows. Integer age also leaves birthday/within-year ambiguity, so compare adjacent eligibility years when scoring sensitivities. @@ -146,7 +157,7 @@ reported source benefit is annual. | `_inferred_aime_lower_bound` | Inclusive conditional AIME interval bound. | | `_inferred_aime_upper_bound` | Exclusive conditional AIME interval bound. | | `_pia_ratio_eligibility_year` | Underlying worker eligibility/computation year E. | -| `_pia_ratio_claiming_adjustment` | Imputed worker claiming adjustment a; one by decomposition convention on the auxiliary composite-ratio route. | +| `_pia_ratio_claiming_adjustment` | Imputed worker claiming adjustment a; zero means unavailable on a positive retirement record, while one is the auxiliary composite-ratio decomposition convention. | | `_pia_ratio_own_record_retired_worker` | Own-record retired-worker indicator. | | `_pia_ratio_missing_claim_age` | Required claiming information is unusable. | | `_pia_ratio_pre_1983` | Benchmark predates supported initial-PIA rounding. | @@ -163,6 +174,15 @@ distribution or percentile calculation, also exclude `_pia_ratio_unmodeled`, `_pia_ratio_missing_claim_age` and `_pia_ratio_pre_1983` records and inspect the proxy and measurement-uncertainty flags. +On a positive retirement record, claiming adjustment zero is an unavailable +diagnostic sentinel until the selected scenario yields a usable adjustment. +Missing claiming data and unsupported pre-1983 records can therefore report +zero while keeping the benefit factor at one. Do not interpret this sentinel +as an actuarial factor, a complete benefit cut, or a default FRA assumption; +filter the associated flags. Disability keeps adjustment one, auxiliary +composite ratios use one by decomposition convention, and a zero-benefit +record keeps the neutral adjustment one. + For exact dime-rounded baseline eligibility PIA z, the interval is `[inverse(z), inverse(z + 0.10))`; the point estimate is its lower endpoint. C2's $2,156.80 baseline PIA therefore gives @@ -184,6 +204,14 @@ This avoids expanding binary representation noise into an additional PIA dime; it does not round every annual source to cents or recover precision already lost when the input was stored. +Named AIME point estimates and interval bounds are stored at Core's float32 +resolution. The adapter rounds them toward positive infinity so an ordinary +cast cannot move a lower endpoint below the baseline inverse. Represented +values can differ from the mathematical endpoint by up to one float32 step; +use the shared pure service for exact interval membership or boundary +calculations. It keeps rational endpoints internally, and the shared +`PIAInverse` exposes exact Fraction bounds as well as high-precision decimals. + A zero benefit produces factor one and zero monetary outputs without division. Negative reported amounts and nonpositive divisors are unsupported and rejected. Factors cannot be negative. With every PIA reform off, factors are exactly one. diff --git a/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py b/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py index d2ce8de4e6e..f0b7d62b6d5 100644 --- a/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py +++ b/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py @@ -34,6 +34,7 @@ ) AMOUNTS = ((18_440.64, 0, 0, 0), (12_000, 6_000, 1_200, 300)) GROUP_ENTITIES = ("household", "tax_unit", "spm_unit", "family", "marital_unit") +GROUP_PLURALS = ("households", "tax_units", "spm_units", "families", "marital_units") def _write_h5(path): @@ -72,8 +73,8 @@ def _situation(): }, } result = {"people": people} - for entity in GROUP_ENTITIES: - result[f"{entity}s"] = {person: {"members": [person]} for person in people} + for plural in GROUP_PLURALS: + result[plural] = {person: {"members": [person]} for person in people} for household in result["households"].values(): household["state_code"] = {YEAR: "TX"} return result @@ -196,14 +197,17 @@ def test_supplied_claim_months_load_from_h5_and_drive_inverse(country_system, tm simulation.calculate("social_security_retirement_claim_age_months", YEAR), [747, 0], ) + # Core carries stored inputs forward through its ordinary dataset contract; + # the input has no formula/uprating, and C rebuilds inferred cohorts by year. np.testing.assert_array_equal( simulation.calculate("social_security_retirement_claim_age_months", YEAR + 1), - [0, 0], + [747, 0], ) adjustment = simulation.calculate( "social_security_retirement_pia_ratio_claiming_adjustment", YEAR ) assert adjustment[0] == pytest.approx(0.7125) + assert adjustment[1] == 0 aime = simulation.calculate("social_security_retirement_inferred_aime", YEAR) assert aime[0] == pytest.approx(5000, abs=0.01) lower = simulation.calculate( @@ -247,3 +251,43 @@ def test_claim_months_are_a_pure_annual_input(country_system): assert variable.default_value == 0 assert not variable.formulas assert not getattr(variable, "uprating", None) + + +def test_named_aime_endpoint_round_trips_after_float32_storage(country_system): + """A nonterminating inverse must survive storage in the Person holder.""" + from decimal import Decimal + from fractions import Fraction + + from policyengine_us.tools.social_security.parameters import bend_points_at + from policyengine_us.tools.social_security.pia import compute_pia + + situation = { + "people": { + "retiree": { + "age": {YEAR: 62}, + "social_security_retirement_reported": {YEAR: 60}, + "social_security_retirement_claim_age_months": {YEAR: 804}, + } + }, + **{plural: {"group": {"members": ["retiree"]}} for plural in GROUP_PLURALS}, + } + simulation = CoreSimulation( + situation=situation, + tax_benefit_system=country_system, + default_input_period=YEAR, + default_calculation_period=YEAR, + ) + aime = simulation.calculate("social_security_retirement_inferred_aime", YEAR)[0] + lower = simulation.calculate( + "social_security_retirement_inferred_aime_lower_bound", YEAR + )[0] + upper = simulation.calculate( + "social_security_retirement_inferred_aime_upper_bound", YEAR + )[0] + assert lower == aime + assert lower < upper + assert Fraction(float(lower)) >= Fraction(50, 9) + assert Fraction(float(upper)) >= Fraction(17, 3) + assert compute_pia( + float(aime), YEAR, bend_points=bend_points_at(country_system.parameters, YEAR) + ) == Decimal("5.0") diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml index e3e181f3e3b..7ce16b8caf8 100644 --- a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml @@ -114,6 +114,7 @@ social_security_retirement_reported: 12000 output: social_security_retirement_pia_ratio_factor: [1, 1, 1] + social_security_retirement_pia_ratio_claiming_adjustment: [0, 0, 0] social_security_retirement_pia_ratio_missing_claim_age: [true, true, true] social_security_retirement: [12000, 12000, 12000] @@ -172,7 +173,6 @@ absolute_error_margin: 0.0001 input: age: 70 - birth_year: 1960 social_security_retirement_entitlement_year: 2025 social_security_retirement_entitlement_month: 6 social_security_retirement_claim_age_months: 780 @@ -182,13 +182,26 @@ social_security_retirement_pia_ratio_eligibility_year: 2022 social_security_retirement_pia_ratio_claiming_adjustment: 0.8666666666666667 +- name: Changing the claiming sensitivity preserves the observed eligibility year + period: 2026 + input: + age: 70 + social_security_retirement_entitlement_year: 2025 + social_security_retirement_entitlement_month: 6 + social_security_retirement_claim_age_months: 780 + social_security_retirement_reported: 12000 + gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [fra] + output: + # The factual age-65 award fixes E independently of the FRA sensitivity. + social_security_retirement_pia_ratio_eligibility_year: 2022 + social_security_retirement_pia_ratio_claiming_adjustment: 1 + - name: A supplied birth year determines the full retirement age cohort period: 2026 absolute_error_margin: 0.01 input: age: 70 birth_year: 1964 - social_security_retirement_worker_eligibility_year: 2026 social_security_retirement_claim_age_months: 744 social_security_retirement_reported: 18117.12 gov.ssa.social_security.pia.bend_points.first: 960 @@ -220,6 +233,25 @@ worker: social_security_retirement_pia_ratio_factor: 1 +- name: A linked eligibility cohort determines FRA before an unsupplied age proxy + period: 2026 + absolute_error_margin: 0.0001 + input: + people: + recipient: + age: 70 + social_security_retirement_worker_id: worker + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 12000 + worker: + age: 70 + social_security_worker_computation_year: 2015 + output: + people: + recipient: + social_security_retirement_pia_ratio_eligibility_year: 2015 + social_security_retirement_pia_ratio_claiming_adjustment: 0.75 + - name: Entitlement dates select the same 62 years 3 months claiming adjustment period: 2026 absolute_error_margin: 0.0001 @@ -238,6 +270,25 @@ social_security_retirement_pia_ratio_claiming_adjustment: 0.7125 social_security_retirement_pia_ratio_missing_claim_age: false +- name: Entitlement dates accept birth year and month without a supplied day + period: 2026 + absolute_error_margin: 0.0001 + input: + age: 62 + birth_year: 1964 + social_security_birth_month: 6 + social_security_retirement_entitlement_year: 2026 + social_security_retirement_entitlement_month: 9 + social_security_retirement_reported: 18440.64 + gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [entitlement_dates] + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_eligibility_year: 2026 + social_security_retirement_pia_ratio_claiming_adjustment: 0.7125 + social_security_retirement_pia_ratio_proxy: true + social_security_retirement_pia_ratio_missing_claim_age: false + - name: Cohort gate leaves an existing beneficiary unchanged period: 2026 input: diff --git a/policyengine_us/tools/social_security/population.py b/policyengine_us/tools/social_security/population.py index 0a4e74ae5e3..023b622c1c8 100644 --- a/policyengine_us/tools/social_security/population.py +++ b/policyengine_us/tools/social_security/population.py @@ -35,6 +35,7 @@ from policyengine_us.tools.social_security_benefit_levers import ( SocialSecurityBenefitSource, ) +from policyengine_us.tools.period_branch import has_input SCENARIOS = { @@ -106,6 +107,7 @@ def _record_basis(person, period, parameters, component): ) age = person("age", period) birth_years = person("birth_year", period) + supplied_birth_year = has_input(person.simulation, "birth_year", period) years, linked, worker_indices = resolve_worker_basis( person, period, component, positive ) @@ -120,6 +122,10 @@ def _record_basis(person, period, parameters, component): policy = parameters.gov.contrib.ssa.benefit_levers.pia_ratio(period.start) if component == "retirement": + for i in np.flatnonzero(positive): + adjustments[i] = Fraction( + 0 + ) # Unknown until the selected scenario is usable. # An entitlement date is a claim date, not age-62 eligibility. In its # absence the scoring year's age rebuilds the cohort independently. # A link without an explicit computation basis is not permission to @@ -143,50 +149,64 @@ def _record_basis(person, period, parameters, component): entitlement_months = person(f"{prefix}_entitlement_month", period) births = [None] * n exact_births = np.zeros(n, dtype=bool) + known_birth_months = np.zeros(n, dtype=bool) for i in np.flatnonzero(positive): j = worker_indices[i] - birth = date(int(birth_years[j]), 1, 2) + known_birth_year = supplied_birth_year and birth_years[j] > 0 + cohort = ( + int(birth_years[j]) + if known_birth_year + else int(years[i]) - 62 + if supplied[i] + else year - int(age[j]) + ) + observed_birth_month = None + if ( + not known_birth_year + and not str(birth_dates[j]) + and entitlement[i] > 0 + and 744 <= months[i] <= 840 + ): + if not 1 <= entitlement_months[i] <= 12: + raise ValueError( + "Social Security entitlement month must lie in [1, 12]" + ) + observed_birth_month = ( + int(entitlement[i]) * 12 + + int(entitlement_months[i]) + - 1 + - int(months[i]) + ) + cohort = observed_birth_month // 12 + birth = date(cohort, 1, 2) exact_birth = False if str(birth_dates[j]): birth = date.fromisoformat(str(birth_dates[j])) exact_birth = True - elif birth_months[j] and birth_days[j]: + known_birth_months[i] = True + elif birth_months[j]: try: - birth = date( - int(birth_years[j]), int(birth_months[j]), int(birth_days[j]) - ) + birth = date(cohort, int(birth_months[j]), int(birth_days[j]) or 2) except ValueError as error: raise ValueError( "Invalid Social Security birth month/day" ) from error - exact_birth = True + exact_birth = bool(birth_days[j]) + known_birth_months[i] = True + elif observed_birth_month is not None: + birth = date(cohort, observed_birth_month % 12 + 1, 2) births[i], exact_births[i] = birth, exact_birth if not supplied[i] and not linked[i]: if exact_birth: years[i] = attainment_month(birth, 744) // 12 - elif entitlement[i] > 0: - # A known claim age plus entitlement month identifies the - # worker's age-62 year without treating claiming as eligibility. - assumed_months = { - "supplied_claim_age": int(months[i]), - "age_62": 744, - "age_65": 780, - "fra": full_retirement_age_months(birth), - "age_70": 840, - "entitlement_dates": 0, - }[scenario] - if 744 <= assumed_months <= 840: - if not 1 <= entitlement_months[i] <= 12: - raise ValueError( - "Social Security entitlement month must lie in [1, 12]" - ) - years[i] = ( - int(entitlement[i]) * 12 - + int(entitlement_months[i]) - - 1 - - assumed_months - + 744 - ) // 12 + elif known_birth_year: + years[i] = birth.year + 62 + proxy[i] = True + elif observed_birth_month is not None: + # Only observed claim age identifies this cohort. Changing + # the sensitivity for a must not change the worker's E. + years[i] = (observed_birth_month + 744) // 12 + proxy[i] = True if not years[i]: years[i] = year - age[i] + 62 proxy[i] = True @@ -196,7 +216,7 @@ def _record_basis(person, period, parameters, component): if not exact_birth: proxy[i] = True if scenario == "entitlement_dates": - if not exact_birth or entitlement[i] <= 0: + if not known_birth_months[i] or entitlement[i] <= 0: missing_claim_age[i] = True continue if not 1 <= entitlement_months[i] <= 12: @@ -352,5 +372,20 @@ def reported_pia_ratio_output(person, period, parameters, component, field): cola_factor=cola_by_year[eligibility], annual_rounding_bound=bound, ) - result[i] = float(getattr(inference, field)) + value = getattr(inference, field) + if field == "factor": + value = np.float32(value) + else: + # A nearest-rounded float32 inverse may lie below its exact + # endpoint and lose a PIA dime. Preserve a representable ceiling. + represented = np.float32(value) + if ( + np.isfinite(represented) + and Decimal.from_float(float(represented)) < value + ): + represented = np.nextafter(represented, np.float32(np.inf)) + value = represented + if not np.isfinite(float(value)): + raise ValueError("PIA ratio output exceeds the supported numeric range") + result[i] = float(value) return result diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_birth_day.py b/policyengine_us/variables/gov/ssa/ss/social_security_birth_day.py index 8b70e458234..35d657ab8f5 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_birth_day.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_birth_day.py @@ -7,4 +7,4 @@ class social_security_birth_day(Variable): entity = Person definition_period = YEAR label = "Social Security birth day" - documentation = "Supplied day of birth for the entitlement_dates population claiming scenario. Zero means unavailable. No formula or uprating. Exact dates preserve SSA age attainment on the day before a birthday, including January 1 births." + documentation = "Optional supplied day of birth for the entitlement_dates population claiming scenario. Zero means unavailable; with a known birth month, day 2 is the explicit approximation and raises the proxy flag. No formula or uprating. Exact dates preserve SSA age attainment on the day before a birthday, including January 1 births." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_birth_month.py b/policyengine_us/variables/gov/ssa/ss/social_security_birth_month.py index c5a9ac05e1a..a5e5a10314c 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_birth_month.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_birth_month.py @@ -7,4 +7,4 @@ class social_security_birth_month(Variable): entity = Person definition_period = YEAR label = "Social Security birth month" - documentation = "Supplied birth month, 1 through 12, for the entitlement_dates population claiming scenario. Zero means unavailable. No formula or uprating. Used with the birth year and supplied birth day; an exact ss_calculator_birth_date can supply the complete date instead." + documentation = "Supplied birth month, 1 through 12, for the entitlement_dates population claiming scenario. Zero means unavailable. No formula or uprating. Used with the birth year; a missing birth day uses the explicit day-2 convention and raises the proxy flag. An exact ss_calculator_birth_date can supply the complete date instead." diff --git a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_claiming_adjustment.py b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_claiming_adjustment.py index ac3c16279b0..95cb42328c0 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_claiming_adjustment.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_claiming_adjustment.py @@ -8,7 +8,7 @@ class social_security_retirement_pia_ratio_claiming_adjustment(Variable): definition_period = YEAR unit = "/1" label = "Social Security retirement pia-ratio claiming adjustment" - documentation = "Worker claiming adjustment used in the baseline inversion under the explicitly selected scenario. Disability uses one. An auxiliary composite worker-benefit ratio already includes its claiming adjustment and is not divided twice." + documentation = "Worker claiming adjustment from the explicitly selected scenario. Zero is an unavailable diagnostic sentinel for a positive retirement record until a usable adjustment is computed, including missing claiming data or an unsupported pre-1983 benchmark; inspect the flags rather than interpreting it as FRA or a benefit cut. A zero-benefit record keeps neutral adjustment one." def formula(person, period, parameters): return reported_pia_ratio_output( From 2a0713526c466b7ec0ebc6fd23ed5bdb15263d5f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 27 Oct 2026 06:26:58 -0700 Subject: [PATCH 70/80] Flag inferred birth years with partially supplied birth dates --- docs/usage/social-security-pia-ratio.md | 4 +++- ...ial_security_retirement_pia_ratio_factor.yaml | 16 ++++++++++++++++ .../tools/social_security/population.py | 4 ++++ 3 files changed, 23 insertions(+), 1 deletion(-) diff --git a/docs/usage/social-security-pia-ratio.md b/docs/usage/social-security-pia-ratio.md index d0f5bd7403e..0cf4c44e887 100644 --- a/docs/usage/social-security-pia-ratio.md +++ b/docs/usage/social-security-pia-ratio.md @@ -69,7 +69,9 @@ calculation uses the documented non-January-1 approximation and raises `birth_year` take precedence over a year derived from current age. When birth information is unavailable but a worker eligibility year is supplied, E minus 62 can supply the approximate birth cohort for FRA and credit rates; this does -not change the supplied eligibility year. A claiming scenario is substantive: +not change the supplied eligibility year. Supplied month/day without an +observed birth year or complete birth date still raises the proxy flag. +A claiming scenario is substantive: in C2, replacing the correct age-62 adjustment of 0.70 with 1 infers AIME $2,978 instead of $5,000 and changes the annual gain from $1,165.248 to $1,664.64 before rounding, an diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml index 7ce16b8caf8..a24b5032f61 100644 --- a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml @@ -289,6 +289,22 @@ social_security_retirement_pia_ratio_proxy: true social_security_retirement_pia_ratio_missing_claim_age: false +- name: A supplied birth month and day still flag an age derived birth year + period: 2026 + absolute_error_margin: 0.0001 + input: + age: 62 + social_security_birth_month: 6 + social_security_birth_day: 15 + social_security_retirement_claim_age_months: 747 + social_security_retirement_reported: 18440.64 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_eligibility_year: 2026 + social_security_retirement_pia_ratio_claiming_adjustment: 0.7125 + social_security_retirement_pia_ratio_proxy: true + - name: Cohort gate leaves an existing beneficiary unchanged period: 2026 input: diff --git a/policyengine_us/tools/social_security/population.py b/policyengine_us/tools/social_security/population.py index 023b622c1c8..e929b8229ee 100644 --- a/policyengine_us/tools/social_security/population.py +++ b/policyengine_us/tools/social_security/population.py @@ -153,6 +153,10 @@ def _record_basis(person, period, parameters, component): for i in np.flatnonzero(positive): j = worker_indices[i] known_birth_year = supplied_birth_year and birth_years[j] > 0 + if not known_birth_year and not str(birth_dates[j]): + # A supplied month/day does not make an age-, eligibility- or + # claim-derived birth year an observed complete date. + proxy[i] = True cohort = ( int(birth_years[j]) if known_birth_year From e09c1fa3a07c2ec9b477ab8b5e7242d1534d5dcc Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 27 Oct 2026 06:41:56 -0700 Subject: [PATCH 71/80] Preserve stacked COLA behavior when the PIA lever is off --- docs/usage/social-security-pia-ratio.md | 2 +- .../tools/social_security_benefit_levers.py | 31 +++++++++++++------ 2 files changed, 23 insertions(+), 10 deletions(-) diff --git a/docs/usage/social-security-pia-ratio.md b/docs/usage/social-security-pia-ratio.md index 0cf4c44e887..9d0ae2e6a36 100644 --- a/docs/usage/social-security-pia-ratio.md +++ b/docs/usage/social-security-pia-ratio.md @@ -89,7 +89,7 @@ basis; conflicting record years are unsupported. For own retirement, a supplied worker eligibility year takes precedence. Without one, exact birth information or an explicitly supplied `birth_year` identifies the age-62 eligibility year before the current-age proxy; known entitlement -dates and the selected claim age can also identify that year. If none of these +dates and a supplied observed claim age can also identify that year. If none of these supplies a basis, each scoring year's proxy is `E = t - age_t + 62`. Rebuild it independently for every annual cross-section. A row whose age stays 70 and has no supplied basis has E=2018 in 2026 and E=2027 in 2035; diff --git a/policyengine_us/tools/social_security_benefit_levers.py b/policyengine_us/tools/social_security_benefit_levers.py index fc2484063f9..fb86fa14647 100644 --- a/policyengine_us/tools/social_security_benefit_levers.py +++ b/policyengine_us/tools/social_security_benefit_levers.py @@ -9,7 +9,6 @@ import numpy as np from policyengine_core.model_api import Enum -from policyengine_us.tools.social_security.records import resolve_worker_basis from policyengine_us.tools.ssa_cola_history import validate_cola_rounding_rule @@ -162,11 +161,29 @@ def reported_joint_cola_factor(person, period, parameters, component): 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. - worker_year, linked, _ = resolve_worker_basis(person, period, component, reported) - if np.any(reported & linked & (worker_year == 0) & (first_year <= 0)): - raise ValueError("Linked worker requires an explicit eligibility year") + 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": @@ -192,11 +209,7 @@ def reported_joint_cola_factor(person, period, parameters, component): if component == "retirement" else np.where(entitlement > 0, historical_first, proxy) ) - inferred = np.where( - worker_year > 0, - worker_year + (worker_year >= 1983), - recipient_origin, - ) + 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: From df07a1ffbec745da4b1ad50c319e34b6d87545fb Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 27 Oct 2026 07:01:20 -0700 Subject: [PATCH 72/80] Test reported retirement cohort and claiming inputs without a country load --- .../unit/test_ss_population_retirement.py | 165 ++++++++++++++++++ 1 file changed, 165 insertions(+) create mode 100644 policyengine_us/tests/unit/test_ss_population_retirement.py diff --git a/policyengine_us/tests/unit/test_ss_population_retirement.py b/policyengine_us/tests/unit/test_ss_population_retirement.py new file mode 100644 index 00000000000..a22c69b87f9 --- /dev/null +++ b/policyengine_us/tests/unit/test_ss_population_retirement.py @@ -0,0 +1,165 @@ +"""Retirement population record resolution without constructing a US model. + +Only Person input arrays, stored-input metadata and a selector snapshot are +supplied here. The production adapter resolves eligibility and calls the real +shared cohort/claiming rules; no formula, parameter tree or simulation is copied. +""" + +from fractions import Fraction +from types import SimpleNamespace + +import numpy as np +import pytest +from policyengine_core.periods import period + +from policyengine_us.tools.social_security.population import _record_basis +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, +) + + +PERIOD = period(2026) + + +class PersonInputs: + ids = ("retiree",) + + def __init__(self, **inputs): + age = inputs.get("age", 62) + values = { + "age": age, + "birth_year": 2026 - age, + "social_security_retirement_reported": 12000, + "social_security_retirement_source": SocialSecurityBenefitSource.reported, + "social_security_retirement_worker_eligibility_year": 0, + "social_security_retirement_worker_id": "", + "social_security_worker_computation_year": 0, + "social_security_retirement_entitlement_year": 0, + "social_security_retirement_entitlement_month": 1, + "social_security_retirement_claim_age_months": 744, + "social_security_birth_month": 0, + "social_security_birth_day": 0, + "ss_calculator_birth_date": "", + "social_security_survivors_reported": 0, + "social_security_dependents_reported": 0, + "social_security_disability_reported": 0, + **inputs, + } + self.inputs = {name: np.asarray([value]) for name, value in values.items()} + # Exercise the actual has_input helper, including the distinction + # between supplied birth year and the ordinary age-derived formula. + self.simulation = SimpleNamespace( + tax_benefit_system=SimpleNamespace( + variables={"birth_year": SimpleNamespace(definition_period="year")} + ), + _get_visible_branch_names=lambda: ("default",), + _user_input_keys={("birth_year", "default", PERIOD)} + if "birth_year" in inputs + else set(), + ) + + def __call__(self, name, requested_period): + assert requested_period == PERIOD + return self.inputs[name] + + +def resolve(*, scenario="supplied_claim_age", **inputs): + parameters = SimpleNamespace( + gov=SimpleNamespace( + contrib=SimpleNamespace( + ssa=SimpleNamespace( + benefit_levers=SimpleNamespace( + pia_ratio=lambda when: SimpleNamespace( + claiming_scenario=[scenario] + ) + ) + ) + ) + ) + ) + return _record_basis(PersonInputs(**inputs), PERIOD, parameters, "retirement") + + +def test_supplied_eligibility_cohort_precedes_an_age_derived_birth_year(): + result = resolve( + age=70, + social_security_retirement_worker_eligibility_year=2015, + ) + assert result["eligibility_year"][0] == 2015 + assert result["claiming_adjustment"][0] == Fraction(3, 4) + assert not result["missing_claim_age"][0] + + +def test_supplied_birth_year_precedes_the_current_age_proxy(): + result = resolve(age=70, birth_year=1964) + assert result["eligibility_year"][0] == 2026 + assert result["claiming_adjustment"][0] == Fraction(7, 10) + + +@pytest.mark.parametrize( + ("scenario", "adjustment"), + [("supplied_claim_age", Fraction(13, 15)), ("fra", Fraction(1))], +) +def test_observed_claim_identifies_one_cohort_across_sensitivity_scenarios( + scenario, adjustment +): + result = resolve( + scenario=scenario, + age=70, + social_security_retirement_entitlement_year=2025, + social_security_retirement_entitlement_month=6, + social_security_retirement_claim_age_months=780, + ) + assert result["eligibility_year"][0] == 2022 + assert result["claiming_adjustment"][0] == adjustment + assert result["proxy"][0] + + +def test_supplied_birth_month_and_day_do_not_identify_a_missing_birth_year(): + result = resolve( + social_security_birth_month=6, + social_security_birth_day=15, + social_security_retirement_claim_age_months=747, + ) + assert result["eligibility_year"][0] == 2026 + assert result["claiming_adjustment"][0] == Fraction(57, 80) + assert result["proxy"][0] + + +def test_entitlement_dates_accept_birth_month_without_day_as_a_proxy(): + result = resolve( + scenario="entitlement_dates", + birth_year=1964, + social_security_birth_month=6, + social_security_retirement_entitlement_year=2026, + social_security_retirement_entitlement_month=9, + ) + assert result["eligibility_year"][0] == 2026 + assert result["claiming_adjustment"][0] == Fraction(57, 80) + assert not result["missing_claim_age"][0] + assert result["proxy"][0] + + +def test_missing_supplied_claim_age_has_no_actuarial_default(): + result = resolve(social_security_retirement_claim_age_months=0) + assert result["claiming_adjustment"][0] == 0 + assert result["missing_claim_age"][0] + assert result["unmodeled"][0] + + +def test_zero_benefit_bypasses_record_fields_and_parameter_reads(): + def zero_person(name, requested_period): + assert requested_period == PERIOD + return { + "social_security_retirement_reported": np.array([0]), + "social_security_retirement_source": np.array( + [SocialSecurityBenefitSource.reported] + ), + }[name] + + result = _record_basis(zero_person, PERIOD, None, "retirement") + assert result["eligibility_year"][0] == 0 + assert result["claiming_adjustment"][0] == 1 + assert not result["positive"][0] + assert not result["missing_claim_age"][0] + assert not result["unmodeled"][0] From 27bd7fdb7803513aded985e35faec5d5b5f04d2f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 28 Oct 2026 04:18:43 -0400 Subject: [PATCH 73/80] Disable timing deadlines for PIA ratio invariant tests --- policyengine_us/tests/unit/test_ss_pia_ratio.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/policyengine_us/tests/unit/test_ss_pia_ratio.py b/policyengine_us/tests/unit/test_ss_pia_ratio.py index d7eb847b59b..948c5ced647 100644 --- a/policyengine_us/tests/unit/test_ss_pia_ratio.py +++ b/policyengine_us/tests/unit/test_ss_pia_ratio.py @@ -3,7 +3,7 @@ from decimal import Decimal from fractions import Fraction -from hypothesis import given, strategies as st +from hypothesis import given, settings, strategies as st import pytest from policyengine_us.tools.social_security.pia import ( @@ -255,6 +255,8 @@ def test_cohort_and_shared_retroactive_modes_have_distinct_applicability(): ) +# These properties check arithmetic, not elapsed time on a shared test host. +@settings(deadline=None) @given( st.decimals(min_value=".01", max_value="1000000", places=2), st.integers(1983, 2035) ) @@ -262,6 +264,7 @@ def test_every_reform_off_is_exactly_one(benefit, year): assert compute_pia_ratio(benefit, year, bend_points=BENDS).factor == D(1) +@settings(deadline=None) @given(st.integers(1, 50000)) def test_baseline_inverse_roundtrips_and_contains_the_original_aime(aime): pia = compute_pia(aime, 2020, bend_points=BENDS) @@ -272,6 +275,7 @@ def test_baseline_inverse_roundtrips_and_contains_the_original_aime(aime): assert compute_pia(result.aime, 2020, bend_points=BENDS) == pia +@settings(deadline=None) @given(st.integers(1, 50000)) def test_everywhere_higher_scale_has_factor_at_least_one(aime): pia = compute_pia(aime, 2020, bend_points=BENDS) @@ -281,6 +285,7 @@ def test_everywhere_higher_scale_has_factor_at_least_one(aime): assert result.factor >= 1 +@settings(deadline=None) @given(st.integers(1, 50000), st.integers(0, 100)) def test_factor_is_never_negative_for_nonnegative_reform_scale(aime, percent): pia = compute_pia(aime, 2020, bend_points=BENDS) @@ -291,6 +296,7 @@ def test_factor_is_never_negative_for_nonnegative_reform_scale(aime, percent): assert result.factor >= 0 +@settings(deadline=None) @given(st.integers(1983, 2035)) def test_zero_remains_zero_under_an_active_reform(year): result = compute_pia_ratio(0, year, REVIEW_PATH, bend_points=BENDS) From ae2d60bbbb4d24b9717981454674c62e9d527010 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 28 Oct 2026 13:49:17 -0400 Subject: [PATCH 74/80] Keep shared PIA switches in calculator reference model --- .../tests/unit/test_ss_calculator_differential.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/policyengine_us/tests/unit/test_ss_calculator_differential.py b/policyengine_us/tests/unit/test_ss_calculator_differential.py index 3228ede8e5b..6597c4d557e 100644 --- a/policyengine_us/tests/unit/test_ss_calculator_differential.py +++ b/policyengine_us/tests/unit/test_ss_calculator_differential.py @@ -23,10 +23,11 @@ CALCULATOR_PREFIXES = ("ss_calculator_", "ss_calculated_") NEW_PARAMETER_NODES = ( - # 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.section_2", + # Keep gov.contrib.ssa.pia and gov.contrib.congress.s770.section_2: + # reported-benefit PIA ratios read these shared formula switches even + # without the household calculator. Calculator-only variables are still + # removed above, so neutral formula dependencies on them remain detectable. + # Benefit levers and statutory COLA history also remain in the reference. "gov.ssa.social_security.pia.bend_points", "gov.ssa.social_security.cola", ) From bc32567e4908452f63702c7d09f07b4040d0c81e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 28 Oct 2026 13:49:50 -0400 Subject: [PATCH 75/80] Resolve disability proxies from operative PIA formula years --- .../unit/test_ss_population_disability.py | 188 ++++++++++++++++++ .../tools/social_security/population.py | 29 ++- 2 files changed, 200 insertions(+), 17 deletions(-) create mode 100644 policyengine_us/tests/unit/test_ss_population_disability.py diff --git a/policyengine_us/tests/unit/test_ss_population_disability.py b/policyengine_us/tests/unit/test_ss_population_disability.py new file mode 100644 index 00000000000..12d00d77f30 --- /dev/null +++ b/policyengine_us/tests/unit/test_ss_population_disability.py @@ -0,0 +1,188 @@ +"""Disability proxy gates using native parameters without constructing a US model.""" + +from fractions import Fraction +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +from policyengine_core.parameters import ParameterNode +from policyengine_core.periods import period + +from policyengine_us.tools.social_security.population import ( + _disability_proxy_year, + _record_basis, +) +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, +) + + +PARAMETER_ROOT = Path(__file__).resolve().parents[2] / "parameters" +PERIOD = period(2026) + + +@pytest.fixture +def parameters(): + # Load the authored defaults into real Core nodes. Only the three small + # subtrees used by the adapter are needed; no country model is built. + root = ParameterNode(data={}) + for path in ( + "gov.contrib.ssa.pia", + "gov.contrib.congress.s770.section_2", + "gov.contrib.ssa.benefit_levers.pia_ratio", + ): + parent = root + parts = path.split(".") + for name in parts[:-1]: + if name not in parent.children: + full_name = ".".join(filter(None, (parent.name, name))) + parent.add_child(name, ParameterNode(full_name, data={})) + parent = getattr(parent, name) + parent.add_child( + parts[-1], + ParameterNode( + path, + directory_path=str(PARAMETER_ROOT.joinpath(*parts)), + ), + ) + return root + + +class DisabilityInputs: + ids = ("recipient",) + + def __init__(self, **inputs): + values = { + "age": 45, + "birth_year": 1981, + "social_security_disability_reported": 12000, + "social_security_disability_source": SocialSecurityBenefitSource.reported, + "social_security_disability_worker_eligibility_year": 0, + "social_security_disability_worker_id": "", + "social_security_worker_computation_year": 0, + "social_security_disability_entitlement_year": 0, + **inputs, + } + self.inputs = {name: np.asarray([value]) for name, value in values.items()} + self.simulation = SimpleNamespace( + tax_benefit_system=SimpleNamespace( + variables={"birth_year": SimpleNamespace(definition_period="year")} + ), + _get_visible_branch_names=lambda: ("default",), + _user_input_keys=set(), + ) + + def __call__(self, name, requested_period): + assert requested_period == PERIOD + return self.inputs[name] + + +@pytest.mark.parametrize("recompute_existing", [False, True]) +def test_missing_onset_proxy_precedes_factor_start_not_enablement_timestamp( + parameters, recompute_existing +): + pia = parameters.gov.contrib.ssa.pia + # Inline YAML reform values begin in 2000; the formula starts in 2026. + pia.in_effect.update(start="2000-01-01", value=True) + pia.recompute_existing.update(start="2000-01-01", value=recompute_existing) + pia.factors.first.update(start="2000-01-01", value=0.95) + pia.factors.start_year.update(start="2000-01-01", value=2026) + + result = _record_basis(DisabilityInputs(), PERIOD, parameters, "disability") + + assert result["eligibility_year"][0] == 2025 + assert result["claiming_adjustment"][0] == Fraction(1) + assert result["proxy"][0] + assert not result["unmodeled"][0] + + +def test_enabled_identity_path_does_not_supply_a_reform_start(parameters): + parameters.gov.contrib.ssa.pia.in_effect.update(start="2000-01-01", value=True) + + assert _disability_proxy_year(parameters, 2026, "pre_reform") == 2025 + assert _disability_proxy_year(parameters, 2035, "pre_reform") == 2034 + + +def test_bend_point_change_uses_its_effective_cohort(parameters): + pia = parameters.gov.contrib.ssa.pia + pia.in_effect.update(start="2000-01-01", value=True) + pia.bend_point_multipliers.first.update(start="2000-01-01", value=1.15) + pia.effective_year.update(start="2000-01-01", value=2024) + + assert _disability_proxy_year(parameters, 2026, "pre_reform") == 2023 + + +def test_disabled_interval_postpones_the_first_operative_change(parameters): + pia = parameters.gov.contrib.ssa.pia + pia.in_effect.update(start="2000-01-01", value=True) + pia.in_effect.update(start="2024-01-01", value=False) + pia.in_effect.update(start="2026-01-01", value=True) + pia.factors.first.update(start="2000-01-01", value=0.95) + pia.factors.start_year.update(start="2000-01-01", value=2024) + + assert _disability_proxy_year(parameters, 2026, "pre_reform") == 2025 + + +def test_s770_start_is_2026_even_when_the_switch_was_enabled_earlier(parameters): + parameters.gov.contrib.congress.s770.section_2.in_effect.update( + start="2000-01-01", value=True + ) + + assert _disability_proxy_year(parameters, 2025, "pre_reform") == 2024 + assert _disability_proxy_year(parameters, 2030, "pre_reform") == 2025 + + +def test_future_formula_change_does_not_backdate_the_proxy(parameters): + pia = parameters.gov.contrib.ssa.pia + pia.in_effect.update(start="2000-01-01", value=True) + pia.factors.first.update(start="2000-01-01", value=0.95) + pia.factors.start_year.update(start="2000-01-01", value=2032) + + assert _disability_proxy_year(parameters, 2026, "pre_reform") == 2025 + + +def test_indexing_identity_anchor_is_still_the_start_of_its_path(parameters): + pia = parameters.gov.contrib.ssa.pia + pia.in_effect.update(start="2000-01-01", value=True) + pia.full_price_indexing.in_effect.update(start="2000-01-01", value=True) + pia.full_price_indexing.start_year.update(start="2000-01-01", value=2024) + pia.full_price_indexing.real_wage_scale.update(start="2024-01-01", value=1) + + assert _disability_proxy_year(parameters, 2026, "pre_reform") == 2023 + + +def test_current_year_sensitivity_bypasses_formula_path_reads(): + assert _disability_proxy_year(None, 2026, "current_year") == 2026 + + +def test_unknown_disability_proxy_is_rejected(): + with pytest.raises(ValueError, match="Choose pre_reform or current_year"): + _disability_proxy_year(None, 2026, "unknown") + + +def test_supplied_worker_year_precedes_entitlement_and_current_year_proxy(parameters): + policy = parameters.gov.contrib.ssa.benefit_levers.pia_ratio + policy.disability_eligibility_proxy.update( + start="1983-01-01", value=["current_year"] + ) + person = DisabilityInputs( + social_security_disability_worker_eligibility_year=2018, + social_security_disability_entitlement_year=2020, + ) + + result = _record_basis(person, PERIOD, parameters, "disability") + + assert result["eligibility_year"][0] == 2018 + assert result["claiming_adjustment"][0] == Fraction(1) + assert not result["proxy"][0] + + +def test_supplied_entitlement_year_precedes_missing_onset_proxy(parameters): + person = DisabilityInputs(social_security_disability_entitlement_year=2020) + + result = _record_basis(person, PERIOD, parameters, "disability") + + assert result["eligibility_year"][0] == 2020 + assert result["claiming_adjustment"][0] == Fraction(1) + assert not result["proxy"][0] diff --git a/policyengine_us/tools/social_security/population.py b/policyengine_us/tools/social_security/population.py index e929b8229ee..37f66298174 100644 --- a/policyengine_us/tools/social_security/population.py +++ b/policyengine_us/tools/social_security/population.py @@ -64,16 +64,14 @@ def _disability_proxy_year(parameters, year, scenario): if scenario != "pre_reform": raise ValueError("Choose pre_reform or current_year for disability eligibility") # This is an exposure sensitivity, not an observed onset distribution. - # Use the cohort immediately before the first enabled formula change. - starts = [year] - node = parameters.gov.contrib.ssa.pia - for entry in node.in_effect.values_list: - if entry.value and int(entry.instant_str[:4]) <= year: - starts.append(int(entry.instant_str[:4])) - for entry in parameters.gov.contrib.congress.s770.section_2.in_effect.values_list: - if entry.value and int(entry.instant_str[:4]) <= year: - starts.append(max(2026, int(entry.instant_str[:4]))) - return max(1983, min(starts) - 1) + # Enablement can precede separate factor/bend/indexing gates (inline YAML + # reforms enable switches from 2000). Use the shared path to find the first + # operative formula change, rather than an enabled but neutral formula. + for candidate in range(1983, year + 1): + path, mode = formula_path_at(parameters, candidate, candidate) + if has_active_reform(path, candidate, candidate, mode): + return max(1983, candidate - 1) + return max(1983, year - 1) def _record_basis(person, period, parameters, component): @@ -249,14 +247,11 @@ def _record_basis(person, period, parameters, component): ) elif component == "disability": years = np.where(years > 0, years, entitlement) - missing = years <= 0 - years = np.where( - missing, - _disability_proxy_year( + missing = positive & (years <= 0) + if np.any(missing): + years[missing] = _disability_proxy_year( parameters, year, _choice(policy.disability_eligibility_proxy) - ), - years, - ) + ) proxy |= missing else: # A worker link/year and an identified composite a*s are both needed. From 22bcfb497c4235cd0498adf38a09d873a33919e7 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 28 Oct 2026 13:52:35 -0400 Subject: [PATCH 76/80] Bound cached PIA reform models per retirement YAML file --- ..._security_retirement_pia_ratio_factor.yaml | 249 +----------------- ..._retirement_pia_ratio_factor_claiming.yaml | 92 +++++++ ...tirement_pia_ratio_factor_eligibility.yaml | 160 +++++++++++ 3 files changed, 255 insertions(+), 246 deletions(-) create mode 100644 policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor_claiming.yaml create mode 100644 policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor_eligibility.yaml diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml index a24b5032f61..5e1c17fb033 100644 --- a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor.yaml @@ -1,3 +1,6 @@ +# Each distinct parameter override combination retains a full model in Core's cache. +# Group these cases to limit each file to at most five combinations. + - name: Inactive population PIA reform preserves a reported record without ages period: 2026 input: @@ -6,95 +9,6 @@ social_security_retirement_pia_ratio_factor: 1 social_security_retirement: 12000 -- name: C2 lower-endpoint inverse preserves the 15 percent bend fixture - period: 2026 - absolute_error_margin: 0.01 - input: - age: 62 - social_security_retirement_claim_age_months: 744 - social_security_retirement_reported: 18117.12 - gov.ssa.social_security.pia.bend_points.first: 960 - gov.ssa.social_security.pia.bend_points.second: 5785 - gov.contrib.ssa.pia.in_effect: true - gov.contrib.ssa.pia.factors.first: 0.95 - gov.contrib.ssa.pia.bend_point_multipliers.first: 1.15 - output: - social_security_retirement_pia_ratio_eligibility_year: 2026 - social_security_retirement_pia_ratio_claiming_adjustment: 0.70 - social_security_retirement_inferred_aime: 5000 - social_security_retirement_inferred_aime_lower_bound: 5000 - social_security_retirement_inferred_aime_upper_bound: 5000.3125 - social_security_retirement_pia_ratio_factor: 1.0643082344 - social_security_retirement: 19282.20 - social_security_retirement_pia_ratio_own_record_retired_worker: true - social_security_retirement_pia_ratio_measurement_uncertainty: true - social_security_retirement_pia_ratio_unmodeled: false - -- name: FRA is an explicit sensitivity and changes inferred AIME - period: 2026 - absolute_error_margin: 0.01 - input: - age: 62 - social_security_retirement_reported: 18117.12 - gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [fra] - gov.ssa.social_security.pia.bend_points.first: 960 - gov.ssa.social_security.pia.bend_points.second: 5785 - gov.contrib.ssa.pia.in_effect: true - gov.contrib.ssa.pia.factors.first: 0.95 - gov.contrib.ssa.pia.bend_point_multipliers.first: 1.15 - output: - social_security_retirement_pia_ratio_claiming_adjustment: 1 - social_security_retirement_inferred_aime: 2978 - social_security_retirement_pia_ratio_proxy: true - -- name: Age 62 sensitivity identifies its claiming adjustment explicitly - period: 2026 - absolute_error_margin: 0.01 - input: - age: 62 - social_security_retirement_reported: 18117.12 - gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [age_62] - gov.ssa.social_security.pia.bend_points.first: 960 - gov.ssa.social_security.pia.bend_points.second: 5785 - gov.contrib.ssa.pia.in_effect: true - gov.contrib.ssa.pia.factors.first: 0.95 - output: - social_security_retirement_pia_ratio_claiming_adjustment: 0.70 - social_security_retirement_inferred_aime: 5000 - social_security_retirement_pia_ratio_proxy: true - -- name: Age 65 sensitivity retains the 24 month early retirement reduction - period: 2026 - absolute_error_margin: 0.01 - input: - age: 62 - social_security_retirement_reported: 22430.72 - gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [age_65] - gov.ssa.social_security.pia.bend_points.first: 960 - gov.ssa.social_security.pia.bend_points.second: 5785 - gov.contrib.ssa.pia.in_effect: true - gov.contrib.ssa.pia.factors.first: 0.95 - output: - social_security_retirement_pia_ratio_claiming_adjustment: 0.8666666666666667 - social_security_retirement_inferred_aime: 5000 - social_security_retirement_pia_ratio_proxy: true - -- name: Age 70 sensitivity uses the birth cohort delayed retirement credits - period: 2026 - absolute_error_margin: 0.01 - input: - age: 62 - social_security_retirement_reported: 32093.184 - gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [age_70] - gov.ssa.social_security.pia.bend_points.first: 960 - gov.ssa.social_security.pia.bend_points.second: 5785 - gov.contrib.ssa.pia.in_effect: true - gov.contrib.ssa.pia.factors.first: 0.95 - output: - social_security_retirement_pia_ratio_claiming_adjustment: 1.24 - social_security_retirement_inferred_aime: 5000 - social_security_retirement_pia_ratio_proxy: true - - name: Missing or out of range supplied claim ages do not fall back period: 2026 input: @@ -148,163 +62,6 @@ social_security_retirement_pia_ratio_pre_1983_count: 1 social_security_retirement: 12000 -- name: This scoring year rebuilds eligibility from its own age cross section - period: 2035 - input: - age: 70 - social_security_retirement_claim_age_months: 744 - social_security_retirement_reported: 12000 - output: - social_security_retirement_pia_ratio_eligibility_year: 2027 - social_security_retirement_pia_ratio_factor: 1 - -- name: The same age in the 2026 cross section has a different eligibility year - period: 2026 - input: - age: 70 - social_security_retirement_claim_age_months: 744 - social_security_retirement_reported: 12000 - output: - social_security_retirement_pia_ratio_eligibility_year: 2018 - social_security_retirement_pia_ratio_factor: 1 - -- name: Supplied entitlement and claim months override a conflicting age proxy - period: 2026 - absolute_error_margin: 0.0001 - input: - age: 70 - social_security_retirement_entitlement_year: 2025 - social_security_retirement_entitlement_month: 6 - social_security_retirement_claim_age_months: 780 - social_security_retirement_reported: 12000 - output: - # First entitlement at 65 in June 2025 identifies June 2022 eligibility. - social_security_retirement_pia_ratio_eligibility_year: 2022 - social_security_retirement_pia_ratio_claiming_adjustment: 0.8666666666666667 - -- name: Changing the claiming sensitivity preserves the observed eligibility year - period: 2026 - input: - age: 70 - social_security_retirement_entitlement_year: 2025 - social_security_retirement_entitlement_month: 6 - social_security_retirement_claim_age_months: 780 - social_security_retirement_reported: 12000 - gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [fra] - output: - # The factual age-65 award fixes E independently of the FRA sensitivity. - social_security_retirement_pia_ratio_eligibility_year: 2022 - social_security_retirement_pia_ratio_claiming_adjustment: 1 - -- name: A supplied birth year determines the full retirement age cohort - period: 2026 - absolute_error_margin: 0.01 - input: - age: 70 - birth_year: 1964 - social_security_retirement_claim_age_months: 744 - social_security_retirement_reported: 18117.12 - gov.ssa.social_security.pia.bend_points.first: 960 - gov.ssa.social_security.pia.bend_points.second: 5785 - gov.contrib.ssa.pia.in_effect: true - gov.contrib.ssa.pia.factors.first: 0.95 - output: - # The supplied cohort has FRA 67 despite this cross section's age proxy. - social_security_retirement_pia_ratio_eligibility_year: 2026 - social_security_retirement_pia_ratio_claiming_adjustment: 0.70 - social_security_retirement_inferred_aime: 5000 - -- name: A linked worker computation year overrides the recipient age proxy - period: 2026 - input: - people: - recipient: - age: 80 - social_security_retirement_worker_id: worker - social_security_retirement_claim_age_months: 744 - social_security_retirement_reported: 12000 - worker: - age: 62 - social_security_worker_computation_year: 2020 - output: - people: - recipient: - social_security_retirement_pia_ratio_eligibility_year: 2020 - worker: - social_security_retirement_pia_ratio_factor: 1 - -- name: A linked eligibility cohort determines FRA before an unsupplied age proxy - period: 2026 - absolute_error_margin: 0.0001 - input: - people: - recipient: - age: 70 - social_security_retirement_worker_id: worker - social_security_retirement_claim_age_months: 744 - social_security_retirement_reported: 12000 - worker: - age: 70 - social_security_worker_computation_year: 2015 - output: - people: - recipient: - social_security_retirement_pia_ratio_eligibility_year: 2015 - social_security_retirement_pia_ratio_claiming_adjustment: 0.75 - -- name: Entitlement dates select the same 62 years 3 months claiming adjustment - period: 2026 - absolute_error_margin: 0.0001 - input: - age: 62 - social_security_birth_month: 6 - social_security_birth_day: 15 - social_security_retirement_entitlement_year: 2026 - social_security_retirement_entitlement_month: 9 - social_security_retirement_reported: 18440.64 - gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [entitlement_dates] - gov.contrib.ssa.pia.in_effect: true - gov.contrib.ssa.pia.factors.first: 0.95 - output: - social_security_retirement_pia_ratio_eligibility_year: 2026 - social_security_retirement_pia_ratio_claiming_adjustment: 0.7125 - social_security_retirement_pia_ratio_missing_claim_age: false - -- name: Entitlement dates accept birth year and month without a supplied day - period: 2026 - absolute_error_margin: 0.0001 - input: - age: 62 - birth_year: 1964 - social_security_birth_month: 6 - social_security_retirement_entitlement_year: 2026 - social_security_retirement_entitlement_month: 9 - social_security_retirement_reported: 18440.64 - gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [entitlement_dates] - gov.contrib.ssa.pia.in_effect: true - gov.contrib.ssa.pia.factors.first: 0.95 - output: - social_security_retirement_pia_ratio_eligibility_year: 2026 - social_security_retirement_pia_ratio_claiming_adjustment: 0.7125 - social_security_retirement_pia_ratio_proxy: true - social_security_retirement_pia_ratio_missing_claim_age: false - -- name: A supplied birth month and day still flag an age derived birth year - period: 2026 - absolute_error_margin: 0.0001 - input: - age: 62 - social_security_birth_month: 6 - social_security_birth_day: 15 - social_security_retirement_claim_age_months: 747 - social_security_retirement_reported: 18440.64 - gov.contrib.ssa.pia.in_effect: true - gov.contrib.ssa.pia.factors.first: 0.95 - output: - social_security_retirement_pia_ratio_eligibility_year: 2026 - social_security_retirement_pia_ratio_claiming_adjustment: 0.7125 - social_security_retirement_pia_ratio_proxy: true - - name: Cohort gate leaves an existing beneficiary unchanged period: 2026 input: diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor_claiming.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor_claiming.yaml new file mode 100644 index 00000000000..fcb8d4e6931 --- /dev/null +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor_claiming.yaml @@ -0,0 +1,92 @@ +# Each distinct parameter override combination retains a full model in Core's cache. +# Group these cases to limit each file to at most five combinations. + +- name: C2 lower-endpoint inverse preserves the 15 percent bend fixture + period: 2026 + absolute_error_margin: 0.01 + input: + age: 62 + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 18117.12 + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + gov.contrib.ssa.pia.bend_point_multipliers.first: 1.15 + output: + social_security_retirement_pia_ratio_eligibility_year: 2026 + social_security_retirement_pia_ratio_claiming_adjustment: 0.70 + social_security_retirement_inferred_aime: 5000 + social_security_retirement_inferred_aime_lower_bound: 5000 + social_security_retirement_inferred_aime_upper_bound: 5000.3125 + social_security_retirement_pia_ratio_factor: 1.0643082344 + social_security_retirement: 19282.20 + social_security_retirement_pia_ratio_own_record_retired_worker: true + social_security_retirement_pia_ratio_measurement_uncertainty: true + social_security_retirement_pia_ratio_unmodeled: false + +- name: FRA is an explicit sensitivity and changes inferred AIME + period: 2026 + absolute_error_margin: 0.01 + input: + age: 62 + social_security_retirement_reported: 18117.12 + gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [fra] + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + gov.contrib.ssa.pia.bend_point_multipliers.first: 1.15 + output: + social_security_retirement_pia_ratio_claiming_adjustment: 1 + social_security_retirement_inferred_aime: 2978 + social_security_retirement_pia_ratio_proxy: true + +- name: Age 62 sensitivity identifies its claiming adjustment explicitly + period: 2026 + absolute_error_margin: 0.01 + input: + age: 62 + social_security_retirement_reported: 18117.12 + gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [age_62] + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_claiming_adjustment: 0.70 + social_security_retirement_inferred_aime: 5000 + social_security_retirement_pia_ratio_proxy: true + +- name: Age 65 sensitivity retains the 24 month early retirement reduction + period: 2026 + absolute_error_margin: 0.01 + input: + age: 62 + social_security_retirement_reported: 22430.72 + gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [age_65] + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_claiming_adjustment: 0.8666666666666667 + social_security_retirement_inferred_aime: 5000 + social_security_retirement_pia_ratio_proxy: true + +- name: Age 70 sensitivity uses the birth cohort delayed retirement credits + period: 2026 + absolute_error_margin: 0.01 + input: + age: 62 + social_security_retirement_reported: 32093.184 + gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [age_70] + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_claiming_adjustment: 1.24 + social_security_retirement_inferred_aime: 5000 + social_security_retirement_pia_ratio_proxy: true + diff --git a/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor_eligibility.yaml b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor_eligibility.yaml new file mode 100644 index 00000000000..d88fb3b29eb --- /dev/null +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_pia_ratio_factor_eligibility.yaml @@ -0,0 +1,160 @@ +# Each distinct parameter override combination retains a full model in Core's cache. +# Group these cases to limit each file to at most five combinations. + +- name: This scoring year rebuilds eligibility from its own age cross section + period: 2035 + input: + age: 70 + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 12000 + output: + social_security_retirement_pia_ratio_eligibility_year: 2027 + social_security_retirement_pia_ratio_factor: 1 + +- name: The same age in the 2026 cross section has a different eligibility year + period: 2026 + input: + age: 70 + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 12000 + output: + social_security_retirement_pia_ratio_eligibility_year: 2018 + social_security_retirement_pia_ratio_factor: 1 + +- name: Supplied entitlement and claim months override a conflicting age proxy + period: 2026 + absolute_error_margin: 0.0001 + input: + age: 70 + social_security_retirement_entitlement_year: 2025 + social_security_retirement_entitlement_month: 6 + social_security_retirement_claim_age_months: 780 + social_security_retirement_reported: 12000 + output: + # First entitlement at 65 in June 2025 identifies June 2022 eligibility. + social_security_retirement_pia_ratio_eligibility_year: 2022 + social_security_retirement_pia_ratio_claiming_adjustment: 0.8666666666666667 + +- name: Changing the claiming sensitivity preserves the observed eligibility year + period: 2026 + input: + age: 70 + social_security_retirement_entitlement_year: 2025 + social_security_retirement_entitlement_month: 6 + social_security_retirement_claim_age_months: 780 + social_security_retirement_reported: 12000 + gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [fra] + output: + # The factual age-65 award fixes E independently of the FRA sensitivity. + social_security_retirement_pia_ratio_eligibility_year: 2022 + social_security_retirement_pia_ratio_claiming_adjustment: 1 + +- name: A supplied birth year determines the full retirement age cohort + period: 2026 + absolute_error_margin: 0.01 + input: + age: 70 + birth_year: 1964 + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 18117.12 + gov.ssa.social_security.pia.bend_points.first: 960 + gov.ssa.social_security.pia.bend_points.second: 5785 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + # The supplied cohort has FRA 67 despite this cross section's age proxy. + social_security_retirement_pia_ratio_eligibility_year: 2026 + social_security_retirement_pia_ratio_claiming_adjustment: 0.70 + social_security_retirement_inferred_aime: 5000 + +- name: A linked worker computation year overrides the recipient age proxy + period: 2026 + input: + people: + recipient: + age: 80 + social_security_retirement_worker_id: worker + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 12000 + worker: + age: 62 + social_security_worker_computation_year: 2020 + output: + people: + recipient: + social_security_retirement_pia_ratio_eligibility_year: 2020 + worker: + social_security_retirement_pia_ratio_factor: 1 + +- name: A linked eligibility cohort determines FRA before an unsupplied age proxy + period: 2026 + absolute_error_margin: 0.0001 + input: + people: + recipient: + age: 70 + social_security_retirement_worker_id: worker + social_security_retirement_claim_age_months: 744 + social_security_retirement_reported: 12000 + worker: + age: 70 + social_security_worker_computation_year: 2015 + output: + people: + recipient: + social_security_retirement_pia_ratio_eligibility_year: 2015 + social_security_retirement_pia_ratio_claiming_adjustment: 0.75 + +- name: Entitlement dates select the same 62 years 3 months claiming adjustment + period: 2026 + absolute_error_margin: 0.0001 + input: + age: 62 + social_security_birth_month: 6 + social_security_birth_day: 15 + social_security_retirement_entitlement_year: 2026 + social_security_retirement_entitlement_month: 9 + social_security_retirement_reported: 18440.64 + gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [entitlement_dates] + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_eligibility_year: 2026 + social_security_retirement_pia_ratio_claiming_adjustment: 0.7125 + social_security_retirement_pia_ratio_missing_claim_age: false + +- name: Entitlement dates accept birth year and month without a supplied day + period: 2026 + absolute_error_margin: 0.0001 + input: + age: 62 + birth_year: 1964 + social_security_birth_month: 6 + social_security_retirement_entitlement_year: 2026 + social_security_retirement_entitlement_month: 9 + social_security_retirement_reported: 18440.64 + gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario: [entitlement_dates] + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_eligibility_year: 2026 + social_security_retirement_pia_ratio_claiming_adjustment: 0.7125 + social_security_retirement_pia_ratio_proxy: true + social_security_retirement_pia_ratio_missing_claim_age: false + +- name: A supplied birth month and day still flag an age derived birth year + period: 2026 + absolute_error_margin: 0.0001 + input: + age: 62 + social_security_birth_month: 6 + social_security_birth_day: 15 + social_security_retirement_claim_age_months: 747 + social_security_retirement_reported: 18440.64 + gov.contrib.ssa.pia.in_effect: true + gov.contrib.ssa.pia.factors.first: 0.95 + output: + social_security_retirement_pia_ratio_eligibility_year: 2026 + social_security_retirement_pia_ratio_claiming_adjustment: 0.7125 + social_security_retirement_pia_ratio_proxy: true + From aa5d225ed8bfe810b053a88c8674f69ab4828109 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 28 Oct 2026 19:55:46 -0400 Subject: [PATCH 77/80] Flag assumed retirement claiming scenarios as proxies --- .../ss-pia-ratio-claiming-proxy.fixed.md | 1 + docs/usage/social-security-pia-ratio.md | 5 ++- .../unit/test_ss_population_retirement.py | 40 +++++++++++++++++++ .../tools/social_security/population.py | 3 +- 4 files changed, 47 insertions(+), 2 deletions(-) create mode 100644 changelog.d/ss-pia-ratio-claiming-proxy.fixed.md diff --git a/changelog.d/ss-pia-ratio-claiming-proxy.fixed.md b/changelog.d/ss-pia-ratio-claiming-proxy.fixed.md new file mode 100644 index 00000000000..e0008535a22 --- /dev/null +++ b/changelog.d/ss-pia-ratio-claiming-proxy.fixed.md @@ -0,0 +1 @@ +Flag assumed retirement claiming-age sensitivities as proxies even when the birth date is exact. diff --git a/docs/usage/social-security-pia-ratio.md b/docs/usage/social-security-pia-ratio.md index 9d0ae2e6a36..fca7ba33b16 100644 --- a/docs/usage/social-security-pia-ratio.md +++ b/docs/usage/social-security-pia-ratio.md @@ -164,7 +164,7 @@ reported source benefit is annual. | `_pia_ratio_missing_claim_age` | Required claiming information is unusable. | | `_pia_ratio_pre_1983` | Benchmark predates supported initial-PIA rounding. | | `_pia_ratio_unmodeled` | This amount lacks a supported identified inversion. | -| `_pia_ratio_proxy` | An eligibility, birth or claiming approximation was used. | +| `_pia_ratio_proxy` | An eligibility, birth or claiming approximation was used, including every resolved `age_62`, `age_65`, `fra` or `age_70` sensitivity even with an exact birth date. | | `_pia_ratio_measurement_uncertainty` | Annual reported measurement is not an exact eligibility PIA. | Zero monetary outputs on records without an inference are placeholders. @@ -175,6 +175,9 @@ whether claiming data are missing or the benchmark predates 1983. For an AIME distribution or percentile calculation, also exclude `_pia_ratio_unmodeled`, `_pia_ratio_missing_claim_age` and `_pia_ratio_pre_1983` records and inspect the proxy and measurement-uncertainty flags. +For identified-only distributions, also exclude `_pia_ratio_proxy` records; +report their sensitivity results separately. The named AIME outputs include +supported proxy inferences, so this filter must be applied by the consumer. On a positive retirement record, claiming adjustment zero is an unavailable diagnostic sentinel until the selected scenario yields a usable adjustment. diff --git a/policyengine_us/tests/unit/test_ss_population_retirement.py b/policyengine_us/tests/unit/test_ss_population_retirement.py index a22c69b87f9..2f97709e5ef 100644 --- a/policyengine_us/tests/unit/test_ss_population_retirement.py +++ b/policyengine_us/tests/unit/test_ss_population_retirement.py @@ -147,6 +147,46 @@ def test_missing_supplied_claim_age_has_no_actuarial_default(): assert result["unmodeled"][0] +@pytest.mark.parametrize( + ("scenario", "adjustment"), + ( + ("age_62", Fraction(11, 15)), + ("age_65", Fraction(41, 45)), + ("fra", Fraction(1)), + ("age_70", Fraction(97, 75)), + ), +) +def test_imputed_claiming_scenario_is_flagged_with_exact_birth(scenario, adjustment): + # Born in 1956: FRA is 796 months. At 62 the reduction is + # 36/180 + 16/240; at 65 it is 16/180; at 70 the credit is 44/150. + result = resolve( + scenario=scenario, + age=70, + ss_calculator_birth_date="1956-06-15", + social_security_retirement_claim_age_months=0, + ) + assert not result["unmodeled"][0] + assert result["eligibility_year"][0] == 2018 + assert result["claiming_adjustment"][0] == adjustment + assert result["proxy"][0] + + +@pytest.mark.parametrize("scenario", ("supplied_claim_age", "entitlement_dates")) +def test_observed_claiming_sources_keep_exact_birth_identified(scenario): + result = resolve( + scenario=scenario, + age=70, + ss_calculator_birth_date="1956-06-15", + social_security_retirement_claim_age_months=780, + social_security_retirement_entitlement_year=2021, + social_security_retirement_entitlement_month=6, + ) + assert not result["unmodeled"][0] + assert result["eligibility_year"][0] == 2018 + assert result["claiming_adjustment"][0] == Fraction(41, 45) + assert not result["proxy"][0] + + def test_zero_benefit_bypasses_record_fields_and_parameter_reads(): def zero_person(name, requested_period): assert requested_period == PERIOD diff --git a/policyengine_us/tools/social_security/population.py b/policyengine_us/tools/social_security/population.py index 37f66298174..c836118f3c9 100644 --- a/policyengine_us/tools/social_security/population.py +++ b/policyengine_us/tools/social_security/population.py @@ -215,7 +215,8 @@ def _record_basis(person, period, parameters, component): for i in np.flatnonzero(positive & (years >= 1983)): birth = births[i] exact_birth = exact_births[i] - if not exact_birth: + # A sensitivity assumes claiming history even with an exact birth. + if not exact_birth or scenario in ("age_62", "age_65", "fra", "age_70"): proxy[i] = True if scenario == "entitlement_dates": if not known_birth_months[i] or entitlement[i] <= 0: From 4b905d402ab53f416628d060faf1abab0978cd05 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 28 Oct 2026 19:56:52 -0400 Subject: [PATCH 78/80] Cover PIA ratio ingress and factor composition --- .../tests/core/test_ss_pia_ratio_inputs.py | 162 +++++++++++++-- .../unit/test_ss_population_composition.py | 185 ++++++++++++++++++ 2 files changed, 335 insertions(+), 12 deletions(-) create mode 100644 policyengine_us/tests/unit/test_ss_population_composition.py diff --git a/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py b/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py index f0b7d62b6d5..a8b61ee10ee 100644 --- a/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py +++ b/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py @@ -10,9 +10,11 @@ import numpy as np import pytest +import yaml from policyengine_core.data import Dataset from policyengine_core.reforms import Reform from policyengine_core.simulations import Simulation as CoreSimulation +from policyengine_core.simulations import SimulationBuilder YEAR = 2026 @@ -29,10 +31,12 @@ "adjusted_gross_income", "income_tax", "household_net_income", + "ssi", "ssi_unearned_income", "snap_unearned_income", ) AMOUNTS = ((18_440.64, 0, 0, 0), (12_000, 6_000, 1_200, 300)) +LATER_YEAR_AMOUNTS = (12_000, 6_000, 3_000, 1_000) GROUP_ENTITIES = ("household", "tax_unit", "spm_unit", "family", "marital_unit") GROUP_PLURALS = ("households", "tax_units", "spm_units", "families", "marital_units") @@ -110,6 +114,12 @@ def apply(self): return NeutralPIARatioHooks(system) +@pytest.fixture(scope="module") +def neutral_stacked_base(country_system): + # One read-only constant-one reference serves every ingress route. + return _neutral_stacked_base(country_system) + + def _snapshot(simulation): requested = { name: np.asarray(simulation.calculate(name, YEAR)).copy() @@ -126,8 +136,23 @@ def _snapshot(simulation): return requested, cached +def _assert_snapshots_equal(current, reference): + assert current[0].keys() == reference[0].keys() + for name in current[0]: + np.testing.assert_array_equal( + current[0][name], reference[0][name], err_msg=name + ) + assert reference[1].keys() <= current[1].keys() + for name in reference[1]: + np.testing.assert_array_equal( + current[1][name], reference[1][name], err_msg=name + ) + + @pytest.mark.parametrize("ingress", ("reported_situation", "public_h5")) -def test_all_off_matches_neutral_stacked_base(country_system, tmp_path, ingress): +def test_all_off_matches_neutral_stacked_base( + country_system, neutral_stacked_base, tmp_path, ingress +): """C adds no current-law change on either supported initial-ingress route.""" if ingress == "public_h5": path = tmp_path / "reported_benefits_2026.h5" @@ -141,7 +166,7 @@ def test_all_off_matches_neutral_stacked_base(country_system, tmp_path, ingress) person[f"{component}_reported"] = person.pop(component) results = [] - for system in (country_system, _neutral_stacked_base(country_system)): + for system in (country_system, neutral_stacked_base): simulation_inputs = ( {"dataset": Dataset.from_file(path)} if ingress == "public_h5" @@ -161,17 +186,130 @@ def test_all_off_matches_neutral_stacked_base(country_system, tmp_path, ingress) ) results.append((requested, cached)) - current, reference = results - assert current[0].keys() == reference[0].keys() - for name in current[0]: - np.testing.assert_array_equal( - current[0][name], reference[0][name], err_msg=name - ) - assert reference[1].keys() <= current[1].keys() - for name in reference[1]: - np.testing.assert_array_equal( - current[1][name], reference[1][name], err_msg=name + _assert_snapshots_equal(*results) + + +def _later_year_situation(): + people = { + "recipient": { + "age": {2024: 70, YEAR: 72}, + "taxable_interest_income": {YEAR: 30_000}, + }, + "spouse": {"age": {2024: 68, YEAR: 70}}, + } + for component, amount in zip(COMPONENTS, LATER_YEAR_AMOUNTS): + people["recipient"][component] = {2024: amount} + situation = { + "people": people, + **{plural: {"unit": {"members": list(people)}} for plural in GROUP_PLURALS}, + } + situation["households"]["unit"]["state_code"] = {YEAR: "TX"} + return situation + + +def _build_later_year_simulation(system, route): + if route == "public_situation": + from policyengine_us import Simulation + + return Simulation(situation=_later_year_situation(), tax_benefit_system=system) + if route == "yaml_builder": + # Match the public-name YAML contract consumed by Core's test runner. + text = yaml.safe_dump({"input": _later_year_situation()}) + return SimulationBuilder().build_from_entities( + system, yaml.safe_load(text)["input"] ) + simulation = SimulationBuilder().build_default_simulation(system, count=2) + simulation.set_input("age", YEAR, [72, 70]) + simulation.set_input("taxable_interest_income", YEAR, [30_000, 0]) + simulation.set_input("state_code", YEAR, ["TX", "TX"]) + for component, amount in zip(COMPONENTS, LATER_YEAR_AMOUNTS): + simulation.set_input(component, 2024, [amount, 0]) + return simulation + + +@pytest.mark.parametrize( + "route", ("public_situation", "dataset_set_input", "yaml_builder") +) +def test_current_law_public_and_later_year_ingress( + country_system, neutral_stacked_base, route +): + """Public inputs retain current law after migration and 2024-to-2026 uprating.""" + results = [] + for system in (country_system, neutral_stacked_base): + simulation = _build_later_year_simulation(system, route) + for component in COMPONENTS: + np.testing.assert_array_equal( + simulation.calculate(f"{component}_pia_ratio_factor", YEAR), 1 + ) + results.append(_snapshot(simulation)) + _assert_snapshots_equal(*results) + assert results[0][0]["social_security_retirement"][0] > 12_000 + + +def test_household_flags_count_members_once(country_system): + situation = _later_year_situation() + for person in situation["people"].values(): + for component in COMPONENTS: + person.pop(component, None) + person["social_security_retirement_reported"] = {YEAR: 12_000} + simulation = CoreSimulation(situation=situation, tax_benefit_system=country_system) + flags = simulation.calculate( + "social_security_retirement_pia_ratio_missing_claim_age", YEAR + ) + counts = simulation.calculate( + "social_security_retirement_pia_ratio_missing_claim_age_count", YEAR + ) + np.testing.assert_array_equal(flags, [True, True]) + np.testing.assert_array_equal(counts, [2]) + assert np.dot(counts, [3]) == np.dot(flags, [3, 3]) == 6 + + +def test_assumed_claim_age_proxy_outputs_exclude_identified_only(country_system): + """Exact birth dates do not turn an assumed claim age into observed history.""" + people = { + person: { + "age": {YEAR: 70}, + "ss_calculator_birth_date": {"eternity": "1956-06-15"}, + "social_security_retirement_reported": {YEAR: 12_000}, + "social_security_retirement_claim_age_months": {YEAR: 0}, + } + for person in ("first_retiree", "second_retiree") + } + situation = { + "people": people, + **{plural: {"unit": {"members": list(people)}} for plural in GROUP_PLURALS}, + } + reform = Reform.from_dict( + { + "gov.contrib.ssa.benefit_levers.pia_ratio.claiming_scenario": { + YEAR: ["age_65"] + } + } + ) + simulation = CoreSimulation( + situation=situation, + tax_benefit_system=country_system, + reform=reform, + default_input_period=YEAR, + default_calculation_period=YEAR, + ) + prefix = "social_security_retirement_pia_ratio_" + proxy = simulation.calculate(f"{prefix}proxy", YEAR) + np.testing.assert_array_equal(proxy, [True, True]) + np.testing.assert_array_equal( + simulation.calculate(f"{prefix}proxy_count", YEAR), [2] + ) + assert np.all( + simulation.calculate("social_security_retirement_inferred_aime", YEAR) > 0 + ) + identified_only = simulation.calculate(f"{prefix}own_record_retired_worker", YEAR) + np.testing.assert_array_equal(identified_only, [True, True]) + for flag in ("unmodeled", "missing_claim_age", "pre_1983", "proxy"): + values = simulation.calculate(f"{prefix}{flag}", YEAR) + if flag != "proxy": + np.testing.assert_array_equal(values, [False, False]) + identified_only = identified_only & ~values + assert not np.any(identified_only) def test_supplied_claim_months_load_from_h5_and_drive_inverse(country_system, tmp_path): diff --git a/policyengine_us/tests/unit/test_ss_population_composition.py b/policyengine_us/tests/unit/test_ss_population_composition.py new file mode 100644 index 00000000000..a347112a855 --- /dev/null +++ b/policyengine_us/tests/unit/test_ss_population_composition.py @@ -0,0 +1,185 @@ +"""Reported-benefit inference and public-formula factor composition. + +The authored SSA parameter subtrees exercise the production adapter and public +formula without constructing a country system or cloning its parameter tree. +""" + +from decimal import Decimal +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +from policyengine_core.parameters import ParameterNode +from policyengine_core.periods import period + +import policyengine_us +from policyengine_us.tools.social_security.cola import ( + cola_factor, + read_parameter_events, +) +from policyengine_us.tools.social_security.population import reported_pia_ratio_output +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, + reported_joint_cola_factor, +) +from policyengine_us.variables.gov.ssa.ss.social_security_retirement import ( + social_security_retirement, +) + +PERIOD = period(2026) +E = 2025 +OLD_PIA = Decimal("2156.8") +NEW_PIA = Decimal("2295.5") + + +@pytest.fixture +def parameters(): + parameter_root = Path(policyengine_us.__file__).resolve().parent / "parameters" + root = ParameterNode(data={}) + for dotted in ( + "gov.contrib.ssa.pia", + "gov.contrib.congress.s770.section_2", + "gov.contrib.ssa.benefit_levers", + "gov.ssa.cola_history", + "gov.ssa.social_security.cola", + "gov.ssa.social_security.pia", + ): + parent = root + parts = dotted.split(".") + for part in parts[:-1]: + if part not in parent.children: + full_name = ".".join(filter(None, (parent.name, part))) + parent.add_child(part, ParameterNode(full_name, data={})) + parent = getattr(parent, part) + parent.add_child( + parts[-1], + ParameterNode(dotted, directory_path=str(parameter_root.joinpath(*parts))), + ) + # Fix the C2 initial-eligibility bends without cloning the whole country. + root.gov.ssa.social_security.pia.bend_points.first.update( + start="2025-01-01", value=960 + ) + root.gov.ssa.social_security.pia.bend_points.second.update( + start="2025-01-01", value=5785 + ) + pia = root.gov.contrib.ssa.pia + pia.in_effect.update(start="2026-01-01", value=True) + pia.recompute_existing.update(start="2026-01-01", value=True) + pia.factors.first.update(start="2026-01-01", value=0.95) + pia.bend_point_multipliers.first.update(start="2026-01-01", value=1.15) + cola = root.gov.contrib.ssa.benefit_levers.cola + cola.pp_change.update(start="2026-01-01", value=0.01) + cola.start_year.update(start="2026-01-01", value=2026) + return root + + +class ReportedRetiree: + ids = ("retiree",) + + def __init__(self, annual_reported): + self.reads = [] + values = { + "age": 63, + "birth_year": 1963, + "social_security_retirement_reported": float(annual_reported), + "social_security_retirement_base": float(annual_reported), + "social_security_retirement_source": SocialSecurityBenefitSource.reported, + "social_security_retirement_worker_eligibility_year": E, + "social_security_retirement_worker_id": "", + "social_security_worker_computation_year": 0, + "social_security_retirement_entitlement_year": 2025, + "social_security_retirement_entitlement_month": 1, + "social_security_retirement_first_cola_year": 2026, + "social_security_retirement_first_cola_month": 1, + "social_security_retirement_claim_age_months": 744, + "social_security_birth_month": 0, + "social_security_birth_day": 0, + "ss_calculator_birth_date": "", + "social_security_survivors_reported": 0, + "social_security_dependents_reported": 0, + "social_security_disability_reported": 0, + # Deliberately implausible final outputs make an accidental read + # of an already-reformed benefit/factor easy to detect. + "social_security_retirement": 1_000_000, + "social_security_retirement_joint_cola_factor": 7.0, + "social_security_retirement_pia_ratio_factor": 9.0, + "social_security_across_the_board_factor": 1.08, + "social_security_retirement_payable_factor": 0.80, + } + self.inputs = {name: np.asarray([value]) for name, value in values.items()} + self.simulation = SimpleNamespace( + tax_benefit_system=SimpleNamespace( + variables={"birth_year": SimpleNamespace(definition_period="year")} + ), + _get_visible_branch_names=lambda: ("default",), + _user_input_keys={("birth_year", "default", PERIOD)}, + ) + + def __call__(self, name, requested_period): + assert requested_period == PERIOD + self.reads.append(name) + return self.inputs[name] + + def set(self, name, value): + self.inputs[name] = np.asarray([value]) + + +def test_adapter_uses_baseline_cola_and_public_formula_applies_each_factor_once( + parameters, +): + events = read_parameter_events(parameters, 2026) + baseline_cola = cola_factor(E, "2026-12", events) + assert baseline_cola == Decimal("1.028") + # A claim at 62 with FRA 67 gives .70. The 2025 benchmark has exactly + # the December 2025 / January 2026 statutory COLA in its 2026 level. + annual = Decimal(12) * Decimal(".70") * OLD_PIA * baseline_cola + assert annual == Decimal("18624.39936") + person = ReportedRetiree(annual) + inferred = reported_pia_ratio_output( + person, PERIOD, parameters, "retirement", "aime" + ) + ratio = reported_pia_ratio_output( + person, PERIOD, parameters, "retirement", "factor" + ) + assert inferred[0] == pytest.approx(5000, abs=0.001) + assert ratio[0] == pytest.approx(float(NEW_PIA / OLD_PIA), abs=1e-7) + assert "social_security_retirement" not in person.reads + assert "social_security_retirement_joint_cola_factor" not in person.reads + assert "social_security_retirement_pia_ratio_factor" not in person.reads + + joint = reported_joint_cola_factor(person, PERIOD, parameters, "retirement") + assert joint[0] == pytest.approx(1.038 / 1.028) + assert joint[0] != 1 + person.set("social_security_retirement_joint_cola_factor", joint[0]) + person.set("social_security_retirement_pia_ratio_factor", ratio[0]) + actual = social_security_retirement.formula(person, PERIOD, parameters) + expected = float(annual) * joint[0] * ratio[0] * 1.08 * 0.80 + assert actual[0] == pytest.approx(expected, rel=1e-12) + + # Change both the configured joint-COLA reform and cached reformed values. + # Neither can enter baseline PIA recovery; the gross amount changes only + # when the public formula applies the new joint-COLA factor. + parameters.gov.contrib.ssa.benefit_levers.cola.pp_change.update( + start="2026-01-01", value=0.02 + ) + person.set("social_security_retirement", actual[0]) + person.set("social_security_retirement_joint_cola_factor", 123.0) + person.set("social_security_retirement_pia_ratio_factor", 456.0) + np.testing.assert_array_equal( + reported_pia_ratio_output(person, PERIOD, parameters, "retirement", "aime"), + inferred, + ) + np.testing.assert_array_equal( + reported_pia_ratio_output(person, PERIOD, parameters, "retirement", "factor"), + ratio, + ) + next_joint = reported_joint_cola_factor(person, PERIOD, parameters, "retirement") + assert next_joint[0] == pytest.approx(1.048 / 1.028) + person.set("social_security_retirement_joint_cola_factor", next_joint[0]) + person.set("social_security_retirement_pia_ratio_factor", ratio[0]) + next_actual = social_security_retirement.formula(person, PERIOD, parameters) + assert next_actual[0] == pytest.approx( + float(annual) * next_joint[0] * ratio[0] * 1.08 * 0.80, rel=1e-12 + ) + assert next_actual[0] / actual[0] == pytest.approx(1.048 / 1.038) From e41131b007a8318ddc995d26b838a036c4dedaed Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 28 Oct 2026 20:14:10 -0400 Subject: [PATCH 79/80] Isolate parameter reforms in PIA ratio input tests --- policyengine_us/tests/core/test_ss_pia_ratio_inputs.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py b/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py index a8b61ee10ee..3d8a0340438 100644 --- a/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py +++ b/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py @@ -288,8 +288,9 @@ def test_assumed_claim_age_proxy_outputs_exclude_identified_only(country_system) ) simulation = CoreSimulation( situation=situation, - tax_benefit_system=country_system, - reform=reform, + # Core's reform= applies in place to an explicitly supplied system. + # Construct the Reform wrapper so its parameter copy stays private. + tax_benefit_system=reform(country_system), default_input_period=YEAR, default_calculation_period=YEAR, ) @@ -326,8 +327,7 @@ def test_supplied_claim_months_load_from_h5_and_drive_inverse(country_system, tm ) simulation = CoreSimulation( dataset=Dataset.from_file(path), - tax_benefit_system=country_system, - reform=reform, + tax_benefit_system=reform(country_system), default_input_period=YEAR, default_calculation_period=YEAR, ) From fea0d26ef98a6771e946f621e8edd2bff11a046c Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 28 Oct 2026 20:44:50 -0400 Subject: [PATCH 80/80] Preserve input carry-forward in isolated PIA test systems --- .../tests/core/test_ss_pia_ratio_inputs.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py b/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py index 3d8a0340438..b7a724af9bf 100644 --- a/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py +++ b/policyengine_us/tests/core/test_ss_pia_ratio_inputs.py @@ -92,6 +92,15 @@ def country_system(): return system +def _private_reformed_system(system, reform): + # Reform clones parameters, unlike CoreSimulation(reform=...) on an + # explicit system. Its generic carry-over default shadows PE-US's, so + # retain the country's input contract on the private wrapper as well. + result = reform(system) + result.auto_carry_over_input_variables = system.auto_carry_over_input_variables + return result + + def _neutral_stacked_base(system): """Restore the original A hooks while retaining the complete stacked model.""" from policyengine_us.model_api import Person, Variable, YEAR @@ -111,7 +120,7 @@ def apply(self): ) self.update_variable(variable) - return NeutralPIARatioHooks(system) + return _private_reformed_system(system, NeutralPIARatioHooks) @pytest.fixture(scope="module") @@ -288,9 +297,7 @@ def test_assumed_claim_age_proxy_outputs_exclude_identified_only(country_system) ) simulation = CoreSimulation( situation=situation, - # Core's reform= applies in place to an explicitly supplied system. - # Construct the Reform wrapper so its parameter copy stays private. - tax_benefit_system=reform(country_system), + tax_benefit_system=_private_reformed_system(country_system, reform), default_input_period=YEAR, default_calculation_period=YEAR, ) @@ -327,7 +334,7 @@ def test_supplied_claim_months_load_from_h5_and_drive_inverse(country_system, tm ) simulation = CoreSimulation( dataset=Dataset.from_file(path), - tax_benefit_system=reform(country_system), + tax_benefit_system=_private_reformed_system(country_system, reform), default_input_period=YEAR, default_calculation_period=YEAR, )