diff --git a/.github/workflows/push.yaml b/.github/workflows/push.yaml index abcff5545f5..e69e0a5a1c9 100644 --- a/.github/workflows/push.yaml +++ b/.github/workflows/push.yaml @@ -501,6 +501,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 ad6a88a0b01..e83fc106209 100644 --- a/.github/workflows/python-tests.yaml +++ b/.github/workflows/python-tests.yaml @@ -162,6 +162,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/Makefile b/Makefile index 69c9172f058..c2efd4e5860 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 ss-calculator-engine ss-calculator-differential core policy-parameters policy remaining-heavy remaining +REST_PYTHON_GROUPS := cached-arrays heavy-a heavy-b core-heavy ss-calculator-engine ss-calculator-differential 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 \ @@ -78,6 +78,11 @@ REST_SS_CALCULATOR_ENGINE_TESTS := $(TESTS)/unit/test_ss_calculator_engine.py # Local macOS/Python 3.13.9, four numeric threads: all eight cases passed # with 2.45 GiB peak RSS, below the ~10 GB per-process target above. REST_SS_CALCULATOR_DIFFERENTIAL_TESTS := $(TESTS)/unit/test_ss_calculator_differential.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. @@ -93,7 +98,7 @@ 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_POLICY_PARAMETER_TESTS) $(REST_REMAINING_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 @@ -166,10 +171,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. @@ -322,6 +329,8 @@ test-other-python-rest-ss-calculator-engine: $(call rest_pytest,ss-calculator-engine,$(REST_SS_CALCULATOR_ENGINE_TESTS)) 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: + $(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: @@ -332,8 +341,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-current-law test-microsimulation: - pytest $(TESTS)/microsimulation --maxfail=0 + @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_current_law.py +test-microsimulation-social-security-current-law: + # 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/changelog.d/ss-benefit-levers.added.md b/changelog.d/ss-benefit-levers.added.md new file mode 100644 index 00000000000..162ee4c65c9 --- /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. 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/docs/social-security-benefit-levers.md b/docs/social-security-benefit-levers.md new file mode 100644 index 00000000000..935fcfe4611 --- /dev/null +++ b/docs/social-security-benefit-levers.md @@ -0,0 +1,185 @@ +# Social Security benefit levers + +The benefit lever interface separates reported annual benefits from public +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. + +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 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. + +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 +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` 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 +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 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. +A weighted `Microsimulation` request also needs an explicit weight for that exact +multi-year period, matching Core. + +## 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/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/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/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/system.py b/policyengine_us/system.py index 1b8a070ec6b..ebbcde84794 100644 --- a/policyengine_us/system.py +++ b/policyengine_us/system.py @@ -36,6 +36,10 @@ ) from .tools.default_uprating import add_default_uprating from .tools.period_branch import delete_input_values +from .tools.social_security_inputs import ( + install_social_security_input_compatibility, + 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, @@ -55,6 +59,8 @@ create_spm_provider, ) +install_social_security_input_compatibility() + COUNTRY_DIR = Path(__file__).parent CURRENT_YEAR = 2024 @@ -95,6 +101,7 @@ class CountryTaxBenefitSystem(TaxBenefitSystem): """ variables_dir = COUNTRY_DIR / "variables" + _social_security_input_compatibility = True auto_carry_over_input_variables = True basic_inputs = [ "state_name", @@ -259,6 +266,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 +528,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/code_health/test_make_test_matches_ci.py b/policyengine_us/tests/code_health/test_make_test_matches_ci.py index 24bb5fef2ac..68437fd99b5 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 @@ -111,3 +112,71 @@ 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" + + +@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" 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 f4318cd3342..33b0c24b2f8 100644 --- a/policyengine_us/tests/code_health/test_rest_python_groups.py +++ b/policyengine_us/tests/code_health/test_rest_python_groups.py @@ -952,6 +952,70 @@ def stub_environment(tmp_path, outcomes): return env, calls +@pytest.mark.parametrize("failed_index", [None, 0, 1, 2]) +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" + 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 = [] + 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]) + 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]) + 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 + 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 (microsimulation_result.returncode != 0) == (failed_index in (0, 1)) + assert (input_result.returncode != 0) == (failed_index == 2) + + def scenario_outcomes(scenario, count): outcomes = ["0"] * count if scenario == "one-fails-one-killed": 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/tests/core/_social_security_state_backend.py b/policyengine_us/tests/core/_social_security_state_backend.py new file mode 100644 index 00000000000..369e18abc0f --- /dev/null +++ b/policyengine_us/tests/core/_social_security_state_backend.py @@ -0,0 +1,779 @@ +"""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. 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 json +import logging +import os +import sys +import types +import warnings +from contextlib import ExitStack, nullcontext +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 +from weakref import WeakKeyDictionary, 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, + advance_sequence_past, + 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] +COMPONENTS = ( + "social_security_retirement", + "social_security_disability", + "social_security_survivors", + "social_security_dependents", +) + + +@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.ClassDef, 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 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 Backend: + """Install the unchanged hooks around unchanged Core routes in isolation.""" + + 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: + self.stack.close() + 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 + branch_clone = core_simulation._branch_clone + BranchClone = core_simulation._BranchClone + + 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, + roles=(member,), + flattened_roles=(member,), + ) + + 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 = InMemoryStorage(is_eternal=False) + self._disk_storage = None + self._on_disk_storable = False + self._do_not_store = name in ( + self.simulation.memory_config.variables_to_drop + ) + + def _storage_period(self, period): + return as_period(period) + + def default_array(self): + return np.zeros(1) + + 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] + 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, + entity=person, + definition_period="year", + is_neutralized=False, + value_type=float, + dtype=np.dtype("float64"), + 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") + or name + in ( + "independent_input", + "raw_seed", + "early_derived", + "late_derived", + ) + ), + ) + for name in ( + "independent_input", + "raw_seed", + "early_derived", + "late_derived", + "unrelated_formula", + *COMPONENTS, + *(f"{name}_reported" for name in COMPONENTS), + *(f"{name}_base" 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 + + def parameters(self, *args): + return types.SimpleNamespace(gov=types.SimpleNamespace(abolitions={})) + + system = System() + + class Simulation: + _open_frames = None + 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._input_epoch = 0 + self._store_history = None + self._requested_variables = None + self._calculations_started = 0 + self._calculations_in_flight = 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 = 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( + 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 + 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 subsample(self, *args): + raise NotImplementedError( + "Dataset subsampling is outside this storage backend" + ) + + def calculate(self, name, period): + # 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): + 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 == "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"): + 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): + 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, + log=logging.getLogger(__name__), + commons=commons, + next_sequence_number=next_sequence_number, + _calculation_frames=calculation_frames, + _formula_simulation=formula_simulation, + _branch_clone=branch_clone, + _BranchClone=BranchClone, + _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=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), + Dataset=Dataset, + Simulation=Simulation, + SimulationBuilder=Builder, + _get_person_count=lambda values: 1, + check_type=lambda *args: None, + PeriodMismatchError=ValueError, + _period_from_storage_key=_period_from_storage_key, + copy=copy, + ) + holder_methods = ( + "clone", + "_get_array_from_storage", + "get_array", + "get_known_periods", + "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", + "_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) + 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) + ) + 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. + 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( + simulation_source, + ("_end_order", "_latest_input_key", "_uprating_index_value"), + namespace, + ) + sim_methods = ( + "set_input", + "delete_arrays", + "_get_store_history", + "_get_requested_variables", + "_drop_values_that_may_depend_on", + "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", + "_calculate", + "calculate_add", + "_cache_option_result", + "_calculation_start", + "_cache_result", + "_may_keep", + "_input_set_meanwhile", + "_drop_computed", + ) + 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_number, + 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__, + ) + 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=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", + 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_module = module( + "_social_security_state_hooks", + json=json, + ContextVar=ContextVar, + WeakKeyDictionary=WeakKeyDictionary, + advance_sequence_past=advance_sequence_past, + ref=ref, + dataclass=dataclass, + wraps=wraps, + Path=Path, + next_sequence_number=next_sequence_number, + get_stored_input=namespace["get_stored_input"], + to_period=as_period, + NpEncoder=core_simulation.NpEncoder, + 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, 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 + self.DelegatingSimulation = DelegatingSimulation + self.Builder = Builder + self.Dataset = Dataset + 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 new file mode 100644 index 00000000000..aabc2ebaabb --- /dev/null +++ b/policyengine_us/tests/core/test_social_security_input_state_machine.py @@ -0,0 +1,1089 @@ +"""Search compositions of public and reported Social Security input routes. + +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. +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: +``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, + initialize, + invariant, + precondition, + 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) +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 +class Input: + amount: int + registered: bool + route: str + + +@dataclass +class Model: + # 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 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, + reported, + period, + amount, + registered, + 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) + if reported: + self.explicit[key] = value + else: + self.public[key] = value + 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): + for name in ("public", "explicit", "overrides"): + setattr( + self, + name, + { + key: value + for key, value in getattr(self, name).items() + if value.registered + }, + ) + + 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 target earlier clones, branches, or restored copies.""" + + def __init__(self): + super().__init__() + 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() + if self.backend is not None: + 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.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 + + 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, + 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) + 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 + ) + 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, + cleanup=self.backend.cleans_raw_overlaps, + ) + + @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) + holder = simulation.get_holder(self.name(component, reported)) + holder.put_in_cache( + [amount], as_period(period), simulation.branch_name, derived=False + ) + # 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, + 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) + # 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(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, + 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( + 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) + + @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 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 + # 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, + branch_name="default" if yaml else model.branch_name, + cleanup=self.backend.cleans_raw_overlaps, + ) + 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) + self.prepare_core_input_drop( + simulation, model, simulation.fail_name, period, amount + ) + 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, + cleanup=self.backend.cleans_raw_overlaps, + ) + + @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() + ) + 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) + + @precondition(lambda self: not self.branch_initial_rewrite_done) + @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. + # 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) + self.reported_values_and_provenance_follow_intent() + self.builder_load( + len(self.states) - 1, + component, + False, + period, + replacement, + fresh=False, + yaml=False, + ) + self.branch_initial_rewrite_done = True + + @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.flatten(promote_raw=True) + self.append(restored, copied_model) + + @rule(slot=TARGET, full_history=st.booleans()) + def extract_reload(self, slot, full_history=False): + simulation, model = self.state(slot) + 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()) + 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 == [] + 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) + 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 source view") + 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 == expected_automatic, ( + 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 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.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 + ) + # 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: + continue + 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() 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..3f5328e40d5 --- /dev/null +++ b/policyengine_us/tests/core/test_social_security_inputs.py @@ -0,0 +1,1468 @@ +"""Input compatibility, branch isolation and complete source round trips. + +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 +from types import SimpleNamespace + +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, + SimulationBuilder, +) + +from policyengine_us import Microsimulation, Simulation +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 ( + 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 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) + + +WRAPPER_WRITE_ROUTES = ( + *((wrapper, "simulation") for wrapper in WRAPPERS), + (Simulation, "holder"), +) +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), + array=np.asarray(values), + branch_name=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] + + +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 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) == 0 + + +@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 = 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) == 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,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, write_route, 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]) + 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 + 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 + ) + # 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( + 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("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: + 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]) + 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, name, 2024) == 24_000 + 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("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( + 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)) +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 + # 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 + + +@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}}) + 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"), + ): + 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 + 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_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("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): + # 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) + + +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 + + +@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_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) + source_holder._disk_storage = source_holder.create_disk_storage( + tmp_path, preserve=True + ) + source_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 source_holder._memory_storage.get(stored_period) is None + np.testing.assert_array_equal( + source_holder._disk_storage.get(stored_period), [12_000] + ) + assert reported_holder._memory_storage.get(stored_period) is None + + # 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( + 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 source_holder._memory_storage.get(stored_period) is None + np.testing.assert_array_equal( + source_holder._disk_storage.get(stored_period), [12_000] + ) + source_holder._on_disk_storable = True + simulation.set_input(public, 2024, [24_000]) + np.testing.assert_array_equal( + 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) + 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 + + +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] + 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 + + +@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( + { + "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 = 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. + simulation.set_input(public, 2024, [0]) + assert amount(simulation, public, 2024) == 0 + 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) == 0 + + +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] + reloaded.set_input(public, 2027, [24_000]) + # 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(): + 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 + 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) +@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"] == override + reloaded = Simulation(situation=extracted, tax_benefit_system=system) + assert amount(reloaded, public, 2024) == 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,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, write_route, snapshot +): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + simulation = make_simulation(wrapper, {public: {2024: 12_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": + 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) + 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) + 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(): + 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 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 ( + 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 + ) + + +@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" + 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 + 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 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) == 0 + 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, wrapper +): + public = SOCIAL_SECURITY_COMPONENTS[0] + reported = f"{public}_reported" + multi_year = "year:2024:2" + 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 + + +@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, multi_year, annual_factors +): + public = SOCIAL_SECURITY_COMPONENTS[0] + simulation = make_simulation(wrapper, {public: {multi_year: 24_000}}) + 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]) + 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} + + +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 + + +@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( + 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) + 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(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: + 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]) + + +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) + + +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 + ) + + +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/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/core/test_social_security_real_core.py b/policyengine_us/tests/core/test_social_security_real_core.py new file mode 100644 index 00000000000..b29fb14af6f --- /dev/null +++ b/policyengine_us/tests/core/test_social_security_real_core.py @@ -0,0 +1,987 @@ +"""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 itertools +import pickle +import sys +import types +from pathlib import Path + +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, + 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 +from policyengine_core.warnings import MemoryConfigWarning + + +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 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]) + + +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 + + +@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("mode", ("blacklist", "variables_to_drop")) +@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, source, 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(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(source)._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")) +@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"]) + 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("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] + # 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]) + 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, 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.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 + 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_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")) + 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 + # 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 ( + 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", "--confcutdir", str(Path(__file__).parent)]) + ) 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/tests/microsimulation/test_social_security_current_law.py b/policyengine_us/tests/microsimulation/test_social_security_current_law.py new file mode 100644 index 00000000000..f07988fd6b2 --- /dev/null +++ b/policyengine_us/tests/microsimulation/test_social_security_current_law.py @@ -0,0 +1,1172 @@ +"""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. +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). +""" + +import io +import json +import os +import subprocess +import sys +import tarfile +import tempfile +from pathlib import Path + +import numpy as np +import pytest + +COMPONENTS = ( + "social_security_retirement", + "social_security_disability", + "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", + "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", + "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", +) +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" +# 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) +) + + +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": {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}, + "is_disabled": {2024: index == 1}, + "takes_up_ssi_if_eligible": {2024: True}, + } + if include_benefits: + 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 result["people"]} + for household in result["households"].values(): + 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 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( + simulation.calculate(variable, period) + ).tolist() + for period in periods + for variable in variables + } + + +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): + """Public-column H5, with source values in its scored year.""" + import h5py + + ids = np.array([1, 2]) + columns = {"person_id": ids, "age": [72, 47]} + for entity in ("household", "tax_unit", "spm_unit", "family", "marital_unit"): + 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 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 ( + 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 + + 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 + + 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) + + 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( + {"gov.contrib.ssa.benefit_levers.across_the_board.rate": {"2024": 0}} + ) + if has_levers + 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), + ): + 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: + if has_levers: + raise + result["complete_export_reload"] = _snapshot( + 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) + 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) + ) + 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 + gc.collect() + + # 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. + 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)) + 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 + 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 ( + 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]) + 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) + ) + 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() + + # 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 + 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) + ) + 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() + + # 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 runs 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() + + # 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, + ) + # 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]}}, + 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]) + 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 + ) + del reloaded + 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 + 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") + if label == "nonzero": + # 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( + _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 + # 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_without_empty_holders(simulation, own_dump) + else: + own_dump = legacy_dump + dump_without_empty_holders(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]) + 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 + ) + 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]) + 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 + ) + 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_without_empty_holders(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_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 + ) + 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() + 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) + 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]) + 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 + 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), + ): + 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], + ) + 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) + ) + 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: + 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, + ) + intended_destination = os.environ.get("PE_US_DIFFERENTIAL_LIVE_VALUES") + if intended_destination: + Path(intended_destination).write_text(json.dumps(intended_differences)) + return result + + +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( + ["git", "status", "--porcelain", "--untracked-files=no"], + cwd=checkout, + 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 ( + REVIEWED_MAIN_SHA, + "refs/remotes/origin/main", + "refs/remotes/gh/main", + "refs/heads/main", + ): + exists = subprocess.run( + ["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() + with tarfile.open(fileobj=io.BytesIO(archive)) as stored: + stored.extractall(checkout, filter="data") + return checkout + pytest.skip( + "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" + ) + + +@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 = {} + 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", + "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, + # 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, ( + f"{label} differential probe failed:\n" + 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"], live_values["current"] + + +@pytest.mark.parametrize( + "route", + ( + "constructor_simulation", + "constructor_microsimulation", + "constructor_downstream", + "later_override_simulation", + "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", + "core_entities_builder", + "core_variables_builder", + "core_flat_dict_builder", + "public_set_input", + "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", + "setter_failed_after_loading", + "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", + "clone_parent", + "branch_parent", + "neutral_reform_branch", + "monthly_simulation", + "monthly_microsimulation", + "multi_year_simulation", + "multi_year_microsimulation", + "h5_scored_year", + *COMPOUND_HELPER_ROUTES, + ), +) +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( + current[route][request], + main[route][request], + err_msg=f"Legacy public ingress differs on {route}: {request}", + ) + + +@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 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. + ( + "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") + Path(sys.argv[2]).write_text(json.dumps(_probe())) 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 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 01e350f6715..460ddca406b 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/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..d5e07de6e72 --- /dev/null +++ b/policyengine_us/tests/policy/contrib/ssa/benefit_levers/social_security_retirement_joint_cola_factor.yaml @@ -0,0 +1,141 @@ +- name: A 2027 CPI-E change preserves precision and retirement exposure boundaries + period: 2027 + absolute_error_margin: 0.001 + input: + gov.contrib.ssa.benefit_levers.cola.start_year: 2027 + gov.contrib.ssa.benefit_levers.cola.pp_change: 0.0015 + 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: + 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 + 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: 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 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 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..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 @@ -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 "uprating" not in public.__dict__ def test_reform_wage_indexes_core_tax_threshold_from_2035(): 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 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_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..3cbfc659f8d --- /dev/null +++ b/policyengine_us/tools/social_security_inputs.py @@ -0,0 +1,1332 @@ +"""Read-time Social Security sources without automatic reported inputs. + +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 weakref import WeakKeyDictionary, ref + +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 + + +SOCIAL_SECURITY_COMPONENTS = ( + "social_security_retirement", + "social_security_disability", + "social_security_survivors", + "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_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 +) +# 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 = {} + + +@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. + persons = getattr(simulation, "persons", None) + return getattr(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() + } + + +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. + """ + if holder is None: + return None + period = holder._storage_period(period) + 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 or not storage.has(period, branch): + continue + 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 _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 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 _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, + 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) + ) + 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() + # 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 + for name in SOCIAL_SECURITY_COMPONENTS: + holder = _holder(simulation, name) + if holder is None: + continue + 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. + _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 + # 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) + 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 (reported, str(period)) not in exported_periods + ): + continue + 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, *, 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") 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 _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 _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): + if numbers: + advance_sequence_past(max(numbers.values())) + 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): + # 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) + 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 + 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 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: + # 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. + number = getattr(fast_cache, "_store_numbers", {}).get(key) + # 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 + 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. + + 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) + # 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(): + 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()): + 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 + 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 is None or number >= since: + array = reference() + if array is not None and simulation._fast_cache.get(key) is array: + # A 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) + 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( + "Replace eager compatibility before installing read-time hooks" + ) + + _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, + cache_result=CoreSimulation._cache_result, + delete_arrays=CoreSimulation.delete_arrays, + 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 + ): + 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 + + @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) + 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 + if ( + name not in _REPORTED_COMPONENTS + or not _enabled(simulation) + or holder.variable.is_neutralized + ): + 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() + ) + return value.array + 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"]) + 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 + + @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 + + 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") + # 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, + value, + 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 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: + # 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 + ) + 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)) + ) + 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) + 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) + return result + finally: + _initial_social_security_writes.reset(writes_token) + _initial_social_security_ingress.reset(token) + 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 + ) + 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 + + 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 + + @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: + _subsampling_social_security_sources.reset(token) + + CoreSimulation.subsample = subsample + + original_extract = CoreSimulation.extract_person + + @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 + } + 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 + + CoreSimulation.extract_person = extract_person + + original_dump = simulation_dumper._dump_holder + + @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 + ): + 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 + ) + else: + for name in _REPORTED_COMPONENTS: + path = Path(directory) / name / _LEGACY_DUMP_FILE + if path.exists(): + _restore_automatic_reported_metadata( + simulation, json.loads(path.read_text()), dumped=True + ) + 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 + ) + finally: + _initial_social_security_ingress.reset(token) + + 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, + ) + + install_social_security_period_compatibility() + + +def export_source_inputs(simulation): + """Export all registered source periods and final overrides.""" + 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: + 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(): + 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(): + 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( + { + "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 + ), + }, + cls=NpEncoder, + ) + ) + + +def reload_source_inputs(payload, **simulation_kwargs): + """Reload complete snapshots, including legacy automatic projections.""" + 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]) + metadata = payload.get(_PROVENANCE) + sources = payload["sources"] + if metadata is not None: + sources = {**sources, _PROVENANCE: metadata} + simulation = Simulation(situation=sources, **simulation_kwargs) + for name, periods in payload["overrides"].items(): + for period, values in periods.items(): + simulation.set_input(name, period, values) + _refresh_input_variables(simulation) + return simulation diff --git a/policyengine_us/tools/social_security_periods.py b/policyengine_us/tools/social_security_periods.py new file mode 100644 index 00000000000..ed8a59cc93b --- /dev/null +++ b/policyengine_us/tools/social_security_periods.py @@ -0,0 +1,150 @@ +"""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.simulations import Simulation as CoreSimulation + +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, + SocialSecurityFund, +) +from policyengine_us.tools.social_security_inputs import SOCIAL_SECURITY_COMPONENTS + + +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 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) + 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 + # 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( + "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]): + 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/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() 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) 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..ea4950d9cc7 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_across_the_board_factor.py @@ -0,0 +1,17 @@ +from policyengine_us.model_api import * + + +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): + 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..e12dde9a5fa --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_joint_cola_factor.py @@ -0,0 +1,18 @@ +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 + unit = "/1" + 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..03739259995 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_payable_factor.py @@ -0,0 +1,23 @@ +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 + unit = "/1" + 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..034b50023c9 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_dependents_pia_ratio_factor.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * + + +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." + ) + + 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..bc04b45475b --- /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" + 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..b9987fb0bf1 100644 --- a/policyengine_us/variables/gov/ssa/ss/social_security_disability.py +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability.py @@ -1,4 +1,7 @@ from policyengine_us.model_api import * +from policyengine_us.tools.social_security_benefit_levers import ( + SocialSecurityBenefitSource, +) class social_security_disability(Variable): @@ -7,4 +10,18 @@ class social_security_disability(Variable): definition_period = YEAR label = "Social Security disability benefits (SSDI)" 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..d088d10392b --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_joint_cola_factor.py @@ -0,0 +1,18 @@ +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 + unit = "/1" + 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..947c742323c --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_payable_factor.py @@ -0,0 +1,14 @@ +from policyengine_us.model_api import * + + +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): + 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..e807b247634 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_disability_pia_ratio_factor.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * + + +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." + ) + + 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..c9834a8d7ee --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_payable_factor.py @@ -0,0 +1,16 @@ +from policyengine_us.model_api import * + + +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." + + 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..668e6957e8e --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_joint_cola_factor.py @@ -0,0 +1,18 @@ +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 + unit = "/1" + 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..43f8dcbce80 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_payable_factor.py @@ -0,0 +1,14 @@ +from policyengine_us.model_api import * + + +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): + 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..50375513765 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_retirement_pia_ratio_factor.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * + + +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." + ) + + 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..b8672e8bfc0 --- /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" + 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..b103db77f4c --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_joint_cola_factor.py @@ -0,0 +1,18 @@ +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 + unit = "/1" + 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..44ac4c71a4a --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_payable_factor.py @@ -0,0 +1,14 @@ +from policyengine_us.model_api import * + + +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): + 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..40be5f4d21d --- /dev/null +++ b/policyengine_us/variables/gov/ssa/ss/social_security_survivors_pia_ratio_factor.py @@ -0,0 +1,15 @@ +from policyengine_us.model_api import * + + +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." + ) + + 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..39fe64d8fd0 --- /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" + 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." diff --git a/pyproject.toml b/pyproject.toml index 9c75e2da78e..24e375e7f24 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 165a1d6ed61..efb56829fd1 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" },