diff --git a/docs/runbook.md b/docs/runbook.md index aaac2903..08c733ea 100644 --- a/docs/runbook.md +++ b/docs/runbook.md @@ -54,6 +54,110 @@ uv run python -m policybench.cli reference-outputs \ --output "$RUN_DIR/us/reference_outputs.csv" ``` +## 2b. Reference-Build Gates + +The published US references are not a bare `reference-outputs` run. They come +from a reference system: a fix module that composes the newest policyengine-us +with the pre-freeze publication conventions and the output-definition adapters +(currently `reference_audit/2026-09-28/fixes/latest_final.py`, built by +`reference_audit/2026-09-28/scripts/build_references_latest.py`; see that +directory's README). Before a reference build is published, it must pass every +gate below. The September 29 build applied gates 1 and 2. Gates 3 and 4 follow +the October 5 audit, which found references resting on engine estimates of +inputs no prompt lists. + +1. **Newest release.** The references come from the newest policyengine-us + release when the sweep begins, and the newest release on PyPI must give the + same values before publishing (rule 1 of + `reference_audit/2026-09-28/README.md`). +2. **Reviewed moves.** The reference builder refuses any scored output that + moves without a reviewed action, a listed action the engine does not + reproduce, and any excluded output that moves without a recheck. +3. **Unlisted-input sweep.** Recompute every output with each engine estimate + of an unlisted input replaced by the prompt's literal reading (unlisted + numbers 0, unlisted facts false) or a documented alternative, and list + every output that moves: + + ```bash + uv run policybench unlisted-input-sweep \ + --fix reference_audit/2026-09-28/fixes/latest_final.py \ + --fix-support reference_audit/2026-09-22/fixes/r19_irs_sales_tax_2025.json \ + --run-dir paper/snapshot/20260501/runs/us_full_run_20260612_policyengine_4_16_1_populace \ + --out-dir results/local/unlisted_input_sweep \ + --acknowledged reference_audit/unlisted_input_acknowledged.json \ + --strict + ``` + + `--fix-support` copies a data file the fix module reads from its own + directory (the 2026-09-28 sales tax convention reads the September 22 IRS + table). Each worker process builds the reference system once (about 40 + seconds); a household simulation takes about half a second on an idle + machine and over a second under load. The county reading tries every county + of every household's state, about 9,000 simulations for the 100-household + run (23 minutes on 12 workers on a loaded machine); `--county-states MD` + limits it. `--estimate` and `--scenario` restrict a run for a quick check. + + The command fixes nothing. It exits 1 when the sweep cannot be trusted (and + on a configuration error or a worker crash): + - its recomputed reference differs from any scored published reference + (the wrong reference system, or a changed engine; pass + `--allow-baseline-mismatch` only to preview a new engine); + - a reference row of a swept household was not recomputed; + - a local tax or credit is nonzero in a state output (gate 4); + - a reading's fixed point did not converge. + + With `--strict` it exits 2 when a scored output moves. One kind of move is + reported but not gated. A county can switch on one of the engine's locality + flags (`in_nyc`, `in_san_francisco`, `in_ny_mctd_zone_2` and its other + `in_*` variables), which makes an unlisted household fact true. The prompt's + rule makes such facts false. The sweep reruns that county with the flags held + at their reference values, and a move that this undoes is marked + `prompt_rules_out` (flag in the `localities` column). A move that survives, + such as New York SNAP's utility allowance region, which reads the county + itself, stays `scored`. + + Each scored move needs one of: + - an exclusion record with reason `reference_depends_on_unlisted_input`; + - an entry in the `--acknowledged` file while a ruling is pending: + `{"acknowledged": [{"scenario_id", "variable", "estimate" ("*" for every + reading), "reading" (optional), "status", "note"}]}`. + + `reference_audit/unlisted_input_acknowledged.json` lists the scored moves on + the current references that await Max's rulings (d963, d974). Remove an + entry when its release lands. A move on an excluded output whose record does + not name the estimate's input is `excluded_other_reason`, and the next + release should add a note to that record. + + The out directory holds `report.md` (the gates and every moved output), + `moves.csv` (one row per output and reading that moved, marked `scored`, + `acknowledged`, `prompt_rules_out`, `excluded_same_input` or + `excluded_other_reason`), + `readings.csv` (what each reading set in each household), + `values.csv.gz` (every simulated value), `baseline.csv` (the recomputed + reference against the CSV), `households.json` (overrides, fixed-point + traces and the local taxes checked) and `summary.json` (provenance: + engine versions and the sha256 of the fix files, scenarios, references and + exclusions). + + `policybench/unlisted_input_sweep.py` registers the estimates + (`ESTIMATES`): the state income tax withheld in the federal SALT deduction + (zero, liability and net readings), the 20% local sales tax, the modeled + Medicare Part B premium (zero) and Medicare enrollment (not enrolled), the + county (Allegany for Maryland), weekly hours + for people with none stated, the mortgage origination year, months of SSDI + receipt, and SSI's disability criterion, plus one reading that applies + every literal reading at once. When an audit finds another estimate, + register it there, with a test, before the next build. +4. **Output scope.** The state income tax and state refundable credit outputs + are state-only. policyengine-us 2.15.17 adds Maryland county tax and New + York City tax into its state income tax, and New York City's refundable + credits and San Francisco's Working Families Tax Credit into its state + refundable credits. The sweep adds `policybench/output_scope.py`'s adapter, + which removes each of them from the engine's lists, and checks that each one + it removes is zero for every household, so the adapter changes no published + reference. `tests/test_output_scope.py` fails when a policyengine-us upgrade + adds a local entry the adapter does not name. + ## 3. Run Claude Separately Run Claude models serially. Claude calls need the main-thread wall timeout, and diff --git a/policybench/cli.py b/policybench/cli.py index 8e894490..25c791f8 100644 --- a/policybench/cli.py +++ b/policybench/cli.py @@ -970,6 +970,16 @@ def main(): ), ) + sweep_parser = subparsers.add_parser( + "unlisted-input-sweep", + help="Recompute a frozen run's references with each engine estimate of " + "an unlisted input replaced by its literal or documented reading, and " + "list every output that moves", + ) + from policybench.unlisted_input_sweep import add_arguments as add_sweep_arguments + + add_sweep_arguments(sweep_parser) + args = parser.parse_args() # Enable disk cache for ordinary eval calls. Contract-failure retry/repair @@ -1571,6 +1581,13 @@ def main(): print(f"sha256: {pointer['sha256']}") print(f"Pointer written to {args.pointer_output}") + elif args.command == "unlisted-input-sweep": + from policybench.unlisted_input_sweep import run as run_unlisted_input_sweep + + code = run_unlisted_input_sweep(args) + if code: + sys.exit(code) + else: parser.print_help() sys.exit(1) diff --git a/policybench/output_scope.py b/policybench/output_scope.py new file mode 100644 index 00000000..084312b3 --- /dev/null +++ b/policybench/output_scope.py @@ -0,0 +1,218 @@ +"""Keep local taxes and credits out of the state outputs. + +PolicyBench's state outputs are state-only (``benchmark_specs.json``): + +- ``state_income_tax_before_refundable_credits``: "state individual income tax after + nonrefundable credits and before refundable credits, excluding local income and + payroll taxes"; +- ``state_refundable_credits``: "total refundable state individual income tax + credits". + +Local income taxes have their own output, ``local_income_tax``. + +policyengine-us builds each engine variable of those names by adding the variables +a parameter list names. On 2.15.17 three of those lists carry local entries +(``SCOPED_LISTS``): + +- ``gov.states.household.state_income_tax_before_refundable_credits``: + ``md_local_income_tax_before_refundable_credits`` (Maryland county income tax, + added by upstream #8888; the 2026-09-28 reference system removes it with + ``reference_audit/2026-09-28/fixes/latest_md_local_output_scope.py``) and + ``nyc_income_tax_before_refundable_credits`` (New York City income tax, which + nothing removed); +- ``gov.states.household.state_refundable_credits``: ``nyc_refundable_credits``; +- ``gov.states.ca.tax.income.credits.refundable``: ``ca_sf_wftc``, San Francisco's + Working Families Tax Credit, which reaches ``state_refundable_credits`` through + ``ca_refundable_credits``. + +``reform`` removes every local entry, so the engine variables match the benchmark +outputs. Only the lists change; no formula or amount does. Maryland county tax still +reaches the federal SALT deduction through ``md_withheld_income_tax``'s county +estimate, and NYC tax net of NYC refundable credits (``nyc_income_tax``) is still in +``local_income_tax``. + +Every one of these is zero for every frozen household: ``in_nyc`` is true only in +the five NYC counties and ``in_san_francisco`` only in San Francisco County, and a +household with no county takes its state's alphabetically first county (Albany +County for New York, Alameda County for California). ``scope_violations`` is the +check: it returns each listed local entry that is nonzero in a simulation. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from typing import Any + +STATE_AGGREGATE = "state_income_tax_before_refundable_credits" +STATE_AGGREGATE_PARAMETER = ( + "gov.states.household.state_income_tax_before_refundable_credits" +) +STATE_REFUNDABLE_PARAMETER = "gov.states.household.state_refundable_credits" +CA_REFUNDABLE_PARAMETER = "gov.states.ca.tax.income.credits.refundable" +# Each list, the local entries 2.15.17 puts in it, and state entries it must keep. +SCOPED_LISTS = { + STATE_AGGREGATE_PARAMETER: { + "local": ( + "md_local_income_tax_before_refundable_credits", + "nyc_income_tax_before_refundable_credits", + ), + "keeps": ( + "md_income_tax_before_refundable_credits", + "ny_income_tax_before_refundable_credits", + ), + }, + STATE_REFUNDABLE_PARAMETER: { + "local": ("nyc_refundable_credits",), + "keeps": ("ny_refundable_credits", "ca_refundable_credits"), + }, + CA_REFUNDABLE_PARAMETER: { + "local": ("ca_sf_wftc",), + "keeps": ("ca_eitc",), + }, +} +LOCAL_INCOME_TAX_COMPONENTS = SCOPED_LISTS[STATE_AGGREGATE_PARAMETER]["local"] +LOCAL_COMPONENTS = tuple( + name for spec in SCOPED_LISTS.values() for name in spec["local"] +) +# Spelled as engine variable names so a new local entry is caught by name: any +# list entry containing one of these fragments is local. +LOCAL_NAME_FRAGMENTS = ( + "_local_", + "nyc_", + "_sf_", + "san_francisco", + "philadelphia", + "kansas_city", + "st_louis", + "wilmington", + "denver", + "yonkers", + "multnomah", +) +SCOPE_INSTANT = "2026-01-01" +# Where the lists are rewritten: the benchmark year and the years after it. +SCOPE_PERIOD = "year:2026-01-01:10" +FIX_ID = "policybench_output_scope" +DESCRIPTION = ( + "Exclude Maryland county tax and New York City and San Francisco taxes and " + "credits from the state income tax and state refundable credit outputs " + "(output definition)" +) + + +def is_local_component(name: str) -> bool: + """Whether a list entry names a local (county or city) tax or credit. + + "state_and_local" (as in the SALT variables) is a state entry, not a local one. + """ + stripped = name.replace("state_and_local", "") + return name in LOCAL_COMPONENTS or any( + fragment in stripped for fragment in LOCAL_NAME_FRAGMENTS + ) + + +def scoped_components(components: Iterable[str]) -> list[str]: + """A list without its local entries, order kept.""" + return [name for name in components if not is_local_component(name)] + + +def local_components(components: Iterable[str]) -> list[str]: + """The local entries of a list, order kept.""" + return [name for name in components if is_local_component(name)] + + +def _parameter(parameters, path: str): + node = parameters + for part in path.split("."): + node = getattr(node, part) + return node + + +def listed_local_components( + parameters, path: str = STATE_AGGREGATE_PARAMETER, instant: str = SCOPE_INSTANT +) -> list[str]: + """The local entries one parameter list adds (default: the state income tax).""" + return local_components(_parameter(parameters, path)(instant)) + + +def all_listed_local_components(parameters, instant: str = SCOPE_INSTANT) -> list[str]: + """Every local entry in every scoped list, in ``SCOPED_LISTS`` order.""" + return [ + name + for path in SCOPED_LISTS + for name in listed_local_components(parameters, path, instant) + ] + + +def remove_local_components(parameters): + """``modify_parameters`` callback: drop every local entry from every list.""" + for path, spec in SCOPED_LISTS.items(): + param = _parameter(parameters, path) + current = list(param(SCOPE_INSTANT)) + kept = scoped_components(current) + if kept != current: + param.update(period=SCOPE_PERIOD, value=kept) + after = param(SCOPE_INSTANT) + assert not local_components(after), (path, after) + for name in spec["keeps"]: + assert name in after, (path, name) + return parameters + + +def installed_parameter_list(path: str = STATE_AGGREGATE_PARAMETER) -> list[str]: + """A list parameter of the installed policyengine-us, read from its YAML. + + The latest value is returned. The file is located without importing + policyengine_us and parsed without building a tax-benefit system (about 40 + seconds), so a test can check every engine upgrade for a new local entry. + """ + import importlib.util + from pathlib import Path + + import yaml + + spec = importlib.util.find_spec("policyengine_us") + if spec is None or not spec.submodule_search_locations: + raise ModuleNotFoundError("policyengine_us is not installed") + root = Path(next(iter(spec.submodule_search_locations))) + path_on_disk = root / "parameters" / (path.replace(".", "/") + ".yaml") + # BaseLoader keeps the 0000-01-01 key a string; the default loader rejects it. + data = yaml.load(path_on_disk.read_text(encoding="utf-8"), Loader=yaml.BaseLoader) + values = data["values"] + return list(values[sorted(values)[-1]]) + + +def scope_violations( + simulation: Any, + components: Sequence[str], + period: int | str, + tolerance: float = 0.005, +) -> dict[str, float]: + """Each listed local entry that is nonzero in ``simulation``, with its amount.""" + violations = {} + for name in components: + amount = float(simulation.calculate(name, period).sum()) + if abs(amount) > tolerance: + violations[name] = amount + return violations + + +def _build_reform(): + from policyengine_core.reforms import Reform + + class reform(Reform): + def apply(self): + self.modify_parameters(remove_local_components) + + reform.__module__ = __name__ + return reform + + +def __getattr__(name: str): + # ``reform`` is built on first access, as a fix module's ``reform`` is read + # (getattr(module, "reform")), so importing the helpers stays cheap: the CLI + # imports this module to build its parser. + if name == "reform": + globals()["reform"] = _build_reform() + return globals()["reform"] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/policybench/unlisted_input_sweep.py b/policybench/unlisted_input_sweep.py new file mode 100644 index 00000000..267e5b87 --- /dev/null +++ b/policybench/unlisted_input_sweep.py @@ -0,0 +1,2069 @@ +"""Sweep a frozen run's references against the engine's estimates of unlisted inputs. + +A US reference is a policyengine-us calculation on the facts the prompt lists. Where +a rule needs an input the prompt never lists, the engine supplies one: a default, a +formula estimate, or an assumed county. The prompt tells models to "treat any +unlisted numeric input as 0 and any other unlisted household fact, boolean, or +status input as false" (``policybench.prompts.TASK_PREFACE``). An output whose +reference moves when such an estimate is replaced by that literal reading, or by a +documented alternative reading, rests on an input the prompt never states. That is +the test of the ``reference_depends_on_unlisted_input`` exclusion rule +(``reference_exclusions.json``). + +The audits found each estimate on its own (r14's weekly hours on 2026-09-22, the +SALT withholding estimate and four others on 2026-10-05). ``ESTIMATES`` registers +them. ``run`` recomputes every output of a frozen run on the reference system that +built its references (a fix module, as ``reference_audit/2026-09-28/scripts/ +sweep.py`` defines one), first as published and then under each registered reading, +and reports every output that moves by more than the $1 exact-match tolerance or +flips a flag. Each move is marked scored or already excluded. + +Every simulation runs on one tax-benefit system per worker process: building the +2.15.17 system takes about 40 seconds and a household simulation about half a +second. + +The swept system adds ``policybench.output_scope.reform`` to the fix module, so its +state income tax and state refundable credit outputs exclude local taxes and +credits as the benchmark defines them. The run checks that the adapter changes no +published reference: each local entry it removes must be zero for every household. + +A county reading that switches on one of the engine's locality flags (in_nyc, +in_san_francisco, ...) makes an unlisted household fact true, which the prompt's +rule makes false. The county is simulated again with those flags held at their +reference values; an output whose move that undoes is reported as +``prompt_rules_out`` and does not fail ``--strict``. A move that survives depends +on the county itself and stays scored. + +Usage (``docs/runbook.md``, "Reference-build gates"): + + uv run policybench unlisted-input-sweep \\ + --fix reference_audit/2026-09-28/fixes/latest_final.py \\ + --fix-support reference_audit/2026-09-22/fixes/r19_irs_sales_tax_2025.json \\ + --out-dir results/local/unlisted_input_sweep \\ + --acknowledged reference_audit/unlisted_input_acknowledged.json --strict + +``--run-dir`` defaults to the frozen run, +``paper/snapshot/20260501/runs/us_full_run_20260612_policyengine_4_16_1_populace``. +""" + +from __future__ import annotations + +import copy +import hashlib +import importlib.util +import json +import math +import multiprocessing +import os +import re +import shutil +import sys +import tempfile +import time +from collections.abc import Callable, Iterable, Sequence +from concurrent.futures import ProcessPoolExecutor, as_completed +from dataclasses import dataclass, field +from datetime import datetime, timezone +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from policybench import output_scope + +if TYPE_CHECKING: + import pandas as pd + +TOLERANCE = 1.0 +BASELINE_TOLERANCE = 1e-3 +NOOP_TOLERANCE = 1e-9 +FIXED_POINT_TOLERANCE = 0.005 +MAX_FIXED_POINT_ITERATIONS = 25 +LITERAL = "literal" +ALTERNATIVE = "alternative" +COMBINED_ESTIMATE = "all_literal_readings" +COMBINED_READING = "all_literal" + +# Inputs the references were built under another name. The frozen manifest stores +# partnership_se_income; policyengine-us 2.15.17 calls it +# partnership_self_employment_net_earnings. The 2026-09-28 reference builder +# (reference_audit/2026-09-28/scripts/build_references_latest.py) builds households +# with sweep.py's build_situation, which renames it. +REFERENCE_INPUT_RENAMES = { + "partnership_se_income": "partnership_self_employment_net_earnings" +} + +# Local income taxes whose only route to the federal SALT deduction is the +# withholding estimate. In policyengine-us 2.15.17 the income-tax line of SALT is +# state_withheld_income_tax plus local_income_tax. md_withheld_income_tax carries an +# AGI-based estimate of Maryland county tax (the county's rate times AGI less the +# largest single standard deduction); local_income_tax adds no Maryland entry, and +# the county liability variable feeds no SALT variable. Replacing the withholding +# estimate with what the household paid therefore has to add the county liability +# back. NYC tax already reaches SALT through local_income_tax (as nyc_income_tax, +# net of NYC refundable credits), so it is not added. +SALT_VIA_WITHHOLDING_ONLY = ("md_local_income_tax_before_refundable_credits",) + +DEFAULT_RUN_DIR = Path( + "paper/snapshot/20260501/runs/us_full_run_20260612_policyengine_4_16_1_populace" +) + + +# --------------------------------------------------------------------------- +# Overrides and readings +# --------------------------------------------------------------------------- + +_GROUPS = { + "tax_unit": "tax_units", + "spm_unit": "spm_units", + "household": "households", +} + + +def _canonical(value: Any) -> Any: + if isinstance(value, bool) or isinstance(value, str): + return value + if isinstance(value, int | float): + return float(value) + return value + + +def _same_value(a: Any, b: Any) -> bool: + """Equal and of the same kind: 0.0 and False are different settings.""" + a, b = _canonical(a), _canonical(b) + return type(a) is type(b) and a == b + + +@dataclass(frozen=True) +class Override: + """Input values a reading sets on top of the household as built for the reference. + + ``people`` maps a person's name in the situation to ``{variable: value}``; the + group fields map a variable to its value on the household's one unit. Values are + for the scenario year. + """ + + people: dict[str, dict[str, Any]] = field(default_factory=dict) + tax_unit: dict[str, Any] = field(default_factory=dict) + spm_unit: dict[str, Any] = field(default_factory=dict) + household: dict[str, Any] = field(default_factory=dict) + + def items(self) -> Iterable[tuple[str, str | None, str, Any]]: + """(entity, person or None, variable, value) for every value set.""" + for name in sorted(self.people): + for variable in sorted(self.people[name]): + yield "person", name, variable, self.people[name][variable] + for entity in _GROUPS: + values = getattr(self, entity) + for variable in sorted(values): + yield entity, None, variable, values[variable] + + def is_empty(self) -> bool: + return not any(True for _ in self.items()) + + def as_dict(self) -> dict[str, Any]: + out: dict[str, Any] = {} + if self.people: + out["people"] = { + name: {k: _canonical(v) for k, v in sorted(values.items())} + for name, values in sorted(self.people.items()) + } + for entity in _GROUPS: + values = getattr(self, entity) + if values: + out[entity] = {k: _canonical(v) for k, v in sorted(values.items())} + return out + + def key(self) -> str: + return json.dumps(self.as_dict(), sort_keys=True) + + def describe(self) -> str: + parts = [] + for entity, person, variable, value in self.items(): + where = person if person is not None else entity + shown = f"{value:.2f}" if isinstance(value, float) else str(value) + parts.append(f"{where}.{variable}={shown}") + return "; ".join(parts) + + def merged(self, other: Override) -> Override: + """Both overrides at once; setting one variable to two values is an error.""" + people = {name: dict(values) for name, values in self.people.items()} + for name, values in other.people.items(): + target = people.setdefault(name, {}) + for variable, value in values.items(): + if variable in target and not _same_value(target[variable], value): + raise ValueError(f"conflicting values for {name}.{variable}") + target[variable] = value + groups = {} + for entity in _GROUPS: + mine = dict(getattr(self, entity)) + for variable, value in getattr(other, entity).items(): + if variable in mine and not _same_value(mine[variable], value): + raise ValueError(f"conflicting values for {entity}.{variable}") + mine[variable] = value + groups[entity] = mine + return Override(people=people, **groups) + + def apply(self, situation: dict, period: str) -> dict: + """A copy of ``situation`` with these values set for ``period``.""" + out = copy.deepcopy(situation) + for name, values in self.people.items(): + if name not in out["people"]: + raise KeyError(f"no person {name!r} in the situation") + for variable, value in values.items(): + out["people"][name][variable] = {period: value} + for entity, plural in _GROUPS.items(): + values = getattr(self, entity) + if not values: + continue + units = out[plural] + if len(units) != 1: + raise ValueError(f"expected one {entity}, found {len(units)}") + unit = next(iter(units.values())) + for variable, value in values.items(): + unit[variable] = {period: value} + return out + + +@dataclass +class ReadingPlan: + """One way to fill an unlisted input for one household.""" + + override: Override + variant: str = "" + simulation: Any = None + trace: list[dict[str, float]] = field(default_factory=list) + converged: bool = True + detail: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class Reading: + """A reading of an unlisted input. + + ``kind`` is ``literal`` for the prompt's own rule (unlisted numbers 0, unlisted + facts false) and ``alternative`` for another reading the stated facts support, + used where the engine already applies the literal one. ``plan`` returns the + household's overrides (several for a reading with variants, none where the + reading cannot apply). + """ + + id: str + kind: str + description: str + plan: Callable[[HouseholdContext], list[ReadingPlan]] + + +@dataclass(frozen=True) +class UnlistedEstimate: + """An input the engine fills in that no PolicyBench prompt lists.""" + + id: str + engine_inputs: tuple[str, ...] + entity: str + engine_behavior: str + why_unlisted: str + found_in: str + readings: tuple[Reading, ...] + + def reading(self, reading_id: str) -> Reading: + for reading in self.readings: + if reading.id == reading_id: + return reading + raise KeyError(f"{self.id} has no reading {reading_id!r}") + + def describe(self) -> dict[str, Any]: + return { + "id": self.id, + "engine_inputs": list(self.engine_inputs), + "entity": self.entity, + "engine_behavior": self.engine_behavior, + "why_unlisted": self.why_unlisted, + "found_in": self.found_in, + "readings": [ + {"id": r.id, "kind": r.kind, "description": r.description} + for r in self.readings + ], + } + + +# --------------------------------------------------------------------------- +# Household context +# --------------------------------------------------------------------------- + + +def discover_locality_flags() -> tuple[str, ...]: + """Household flags that put a household in a locality with its own rules. + + policyengine-us defines them as ``in_*`` variables under ``variables/gov/local`` + (in_san_francisco, in_la, ...) and ``variables/household/demographic/geographic`` + (in_nyc, and in_ny_mctd_zone_2 for the Metropolitan Commuter Transportation + District outside the city). Most are formulas of county; a few (in_denver, + in_wilmington, in_yonkers) are inputs that default to false, which no county + reading switches on. They are read from the installed source tree, so no system + is built. The prompt states no locality, and its rule makes every unlisted + household fact false. + """ + spec = importlib.util.find_spec("policyengine_us") + if spec is None or not spec.submodule_search_locations: + return () + root = Path(next(iter(spec.submodule_search_locations))) / "variables" + names = set() + for pattern in ("gov/local/**/in_*.py", "household/demographic/geographic/in_*.py"): + names.update(path.stem for path in root.glob(pattern)) + return tuple(sorted(names)) + + +class Engine: + """The swept tax-benefit system and how to simulate a situation on it.""" + + def __init__( + self, + system: Any, + simulation_class: Any, + aggregate_local_components: Sequence[str] = (), + removed_local_components: Sequence[str] = (), + remaining_local_components: Sequence[str] | None = None, + county_states: frozenset[str] | None = None, + locality_flags: Sequence[str] | None = None, + ): + self.system = system + self.simulation_class = simulation_class + # Local taxes still in the swept state income tax aggregate. + self.aggregate_local_components = tuple(aggregate_local_components) + # Local entries the output-scope adapter removed from the state lists. + self.removed_local_components = tuple(removed_local_components) + # Local entries still in any state list (taxes and credits). + self.remaining_local_components = tuple( + aggregate_local_components + if remaining_local_components is None + else remaining_local_components + ) + self.county_states = county_states + if locality_flags is None: + locality_flags = [ + name for name in discover_locality_flags() if self.has_variable(name) + ] + self.locality_flags = tuple(locality_flags) + + @property + def remaining_local_credits(self) -> tuple[str, ...]: + """Local credits the swept system still counts in state refundable credits.""" + return tuple( + name + for name in self.remaining_local_components + if name not in output_scope.LOCAL_INCOME_TAX_COMPONENTS + and name not in self.aggregate_local_components + ) + + def simulate(self, situation: dict) -> Any: + return self.simulation_class( + tax_benefit_system=self.system, situation=situation + ) + + def has_variable(self, name: str) -> bool: + return name in self.system.variables + + def variable_entity(self, name: str) -> str: + entity = getattr(self.system.variables[name], "entity", None) + return getattr(entity, "key", "household") + + def parameter(self, path: str, instant: str) -> Any: + node = self.system.parameters + for part in path.split("."): + node = getattr(node, part) + return node(instant) + + def counties(self, state: str) -> list[str]: + """Every county the engine knows in ``state``, in the engine's order. + + policyengine-us gives a household with no county its state's first county in + this order (``first_county_in_state``: the alphabetically first value). + """ + enum = self.system.variables["county"].possible_values + members = [ + member + for member in enum + if member.name != "UNKNOWN" and str(member.value).endswith(f", {state}") + ] + return [member.name for member in sorted(members, key=lambda m: str(m.value))] + + +@dataclass +class HouseholdContext: + scenario: Any + situation: dict + variables: list[str] + engine: Engine + baseline: Any = None + simulations: int = 0 + _cache: dict[str, Any] = field(default_factory=dict) + + @property + def year(self) -> int: + return int(self.scenario.year) + + @property + def period(self) -> str: + return str(self.scenario.year) + + @property + def person_names(self) -> list[str]: + return list(self.situation["people"]) + + def total(self, simulation: Any, variable: str) -> float: + return float(simulation.calculate(variable, self.year).sum()) + + def by_person(self, simulation: Any, variable: str) -> dict[str, float]: + values = simulation.calculate(variable, self.year) + return { + name: float(values[index]) for index, name in enumerate(self.person_names) + } + + def _baseline_value(self, entity: str, person: str | None, variable: str) -> Any: + values = self.baseline.calculate(variable, self.year) + if hasattr(values, "decode_to_str"): + values = values.decode_to_str() + index = self.person_names.index(person) if entity == "person" else 0 + value = values[index] + return value if isinstance(value, str) else float(value) + + def is_noop(self, override: Override) -> bool: + """Whether every value ``override`` sets equals the reference's own value. + + Setting an input to the value the engine computes for it leaves every + output unchanged, so such a reading needs no simulation. + """ + for entity, person, variable, value in override.items(): + current = self._baseline_value(entity, person, variable) + if isinstance(current, str) or isinstance(value, str): + if str(current) != str(value): + return False + elif abs(float(current) - float(value)) > NOOP_TOLERANCE: + return False + return True + + def simulate(self, override: Override, *, cache: bool = True) -> Any: + """The household under ``override``. Kept for reuse unless ``cache`` is off + (the county reading's hundreds of variants keep only their outputs).""" + if override.is_empty() or self.is_noop(override): + return self.baseline + key = override.key() + if key in self._cache: + return self._cache[key] + self.simulations += 1 + simulation = self.engine.simulate(override.apply(self.situation, self.period)) + if cache: + self._cache[key] = simulation + return simulation + + def outputs(self, simulation: Any) -> dict[str, float]: + from policybench.ground_truth import ( + _extract_person_value, + _pe_variable_for_output, + ) + + values = {} + for variable in self.variables: + pe_variable = _pe_variable_for_output(variable, self.scenario.country) + values[variable] = float( + _extract_person_value( + simulation.calculate(pe_variable, self.year), + self.scenario, + variable, + ) + ) + return values + + +# --------------------------------------------------------------------------- +# Readings +# --------------------------------------------------------------------------- + + +def _zero_on_tax_unit(variable: str) -> Callable[[HouseholdContext], list[ReadingPlan]]: + def plan(ctx: HouseholdContext) -> list[ReadingPlan]: + return [ + ReadingPlan( + Override(tax_unit={variable: 0.0}), + detail={"engine_value": ctx.total(ctx.baseline, variable)}, + ) + ] + + return plan + + +def _zero_on_people(variable: str) -> Callable[[HouseholdContext], list[ReadingPlan]]: + def plan(ctx: HouseholdContext) -> list[ReadingPlan]: + values = ctx.by_person(ctx.baseline, variable) + people = {name: {variable: 0.0} for name, value in values.items() if value} + if not people: + return [] + return [ReadingPlan(Override(people=people), detail={"engine_value": values})] + + return plan + + +def salt_income_tax_paid( + ctx: HouseholdContext, simulation: Any, *, net: bool = False +) -> float: + """The state income tax a household pays during the year if it pays its liability. + + The engine's state income tax before refundable credits, less any local tax the + swept system's state aggregate still adds, plus the local tax whose only route + to SALT is the withholding estimate (Maryland county tax). ``net`` also + subtracts state refundable credits, less any local credit the swept system still + counts in them. Floored at zero. + + It inherits the engine aggregate's entries: Mississippi's is before + nonrefundable credits (ms_income_tax_before_credits_unit), and Washington's + adds its capital gains and millionaires taxes. + """ + amount = ctx.total(simulation, output_scope.STATE_AGGREGATE) + for name in ctx.engine.aggregate_local_components: + amount -= ctx.total(simulation, name) + for name in SALT_VIA_WITHHOLDING_ONLY: + if ctx.engine.has_variable(name): + amount += ctx.total(simulation, name) + if net: + amount -= ctx.total(simulation, "state_refundable_credits") + for name in ctx.engine.remaining_local_credits: + amount += ctx.total(simulation, name) + return max(0.0, amount) + + +def _withholding_fixed_point( + net: bool, +) -> Callable[[HouseholdContext], list[ReadingPlan]]: + def plan(ctx: HouseholdContext) -> list[ReadingPlan]: + estimate = ctx.total(ctx.baseline, "state_withheld_income_tax") + withheld = salt_income_tax_paid(ctx, ctx.baseline, net=net) + trace = [] + converged = False + override = Override(tax_unit={"state_withheld_income_tax": withheld}) + simulation = ctx.baseline + # A state's tax can read federal tax or the federal SALT deduction, so the + # amount paid is iterated until it reproduces itself. + for _ in range(MAX_FIXED_POINT_ITERATIONS): + override = Override(tax_unit={"state_withheld_income_tax": withheld}) + simulation = ctx.simulate(override) + implied = salt_income_tax_paid(ctx, simulation, net=net) + trace.append({"withheld": withheld, "implied": implied}) + if abs(implied - withheld) < FIXED_POINT_TOLERANCE: + converged = True + break + withheld = implied + return [ + ReadingPlan( + override, + simulation=simulation, + trace=trace, + converged=converged, + detail={ + "engine_value": estimate, + "paid": override.tax_unit["state_withheld_income_tax"], + }, + ) + ] + + return plan + + +def _stated_hours_alternative( + hours: float, +) -> Callable[[HouseholdContext], list[ReadingPlan]]: + target = "weekly_hours_worked_before_lsr" + + def plan(ctx: HouseholdContext) -> list[ReadingPlan]: + # Stated usual hours reach the target through PE_INPUT_ALIASES; a person + # without them keeps the engine default. + people = { + name: {target: hours} + for name, person in ctx.situation["people"].items() + if target not in person + } + if not people: + return [] + return [ReadingPlan(Override(people=people))] + + return plan + + +def _people_flag_alternative( + variable: str, + value: Any, + when: Callable[[HouseholdContext], dict[str, bool]], +) -> Callable[[HouseholdContext], list[ReadingPlan]]: + def plan(ctx: HouseholdContext) -> list[ReadingPlan]: + people = {name: {variable: value} for name, hit in when(ctx).items() if hit} + if not people: + return [] + return [ReadingPlan(Override(people=people))] + + return plan + + +def _receives_ssdi(ctx: HouseholdContext) -> dict[str, bool]: + values = ctx.by_person(ctx.baseline, "social_security_disability") + return {name: value > 0 for name, value in values.items()} + + +def _is_medicare_eligible(ctx: HouseholdContext) -> dict[str, bool]: + values = ctx.by_person(ctx.baseline, "is_medicare_eligible") + return {name: bool(value) for name, value in values.items()} + + +def _is_disabled(ctx: HouseholdContext) -> dict[str, bool]: + values = ctx.by_person(ctx.baseline, "is_disabled") + return {name: bool(value) for name, value in values.items()} + + +def _pre_tcja_mortgage(ctx: HouseholdContext) -> list[ReadingPlan]: + # deductible_mortgage_interest_tax_unit is the only variable that reads the + # origination years; it is interest times a deductible share, min(1, capped + # balance / balance), which raising a cap cannot lower. Where the share is + # already 1 (deductible equals interest), every output is unchanged, so no + # simulation is needed. Testing the share, not non-deductible interest, keeps + # this exact for negative interest too. + interest = ctx.total(ctx.baseline, "home_mortgage_interest_tax_unit") + deductible = ctx.total(ctx.baseline, "deductible_mortgage_interest_tax_unit") + capped = interest - deductible + if abs(capped) <= NOOP_TOLERANCE: + return [] + year = int( + ctx.engine.parameter( + "gov.irs.deductions.itemized.interest.mortgage.pre_tcja_origination_year", + f"{ctx.period}-01-01", + ) + ) + values = {} + for prefix in ("first", "second"): + if ctx.total(ctx.baseline, f"{prefix}_home_mortgage_balance") > 0: + values[f"{prefix}_home_mortgage_origination_year"] = year + if not values: + return [] + return [ + ReadingPlan( + Override(tax_unit=values), + detail={"capped_interest": capped, "origination_year": year}, + ) + ] + + +def _each_county(ctx: HouseholdContext) -> list[ReadingPlan]: + states = ctx.engine.county_states + if states is not None and ctx.scenario.state not in states: + return [] + return [ + ReadingPlan(Override(household={"county": county}), variant=county) + for county in ctx.engine.counties(ctx.scenario.state) + ] + + +ESTIMATES: tuple[UnlistedEstimate, ...] = ( + UnlistedEstimate( + id="state_withheld_income_tax", + engine_inputs=("state_withheld_income_tax",), + entity="tax_unit", + engine_behavior=( + "The income-tax line of the federal SALT deduction is " + "state_withheld_income_tax plus local_income_tax, and it counts when it " + "exceeds general sales tax. state_withheld_income_tax adds per-state " + "*_withheld_income_tax formulas on each person's federal AGI " + "(parameters/gov/states/household/state_withheld_income_tax.yaml); " + "Maryland's includes a county estimate. CO and MO formulas read it, and " + "HI, VA, SC, NM, ID, UT and AZ formulas read SALT." + ), + why_unlisted=( + "No prompt states state income tax withheld or paid during the year " + "(26 U.S.C. 164(a)(3), (b)(5))." + ), + found_in="reference_audit 2026-10-05 (PR #191)", + readings=( + Reading( + id="zero", + kind=LITERAL, + description=( + "No state income tax was withheld or paid (unlisted numeric " + "input = 0); SALT takes local income tax or general sales tax." + ), + plan=_zero_on_tax_unit("state_withheld_income_tax"), + ), + Reading( + id="liability", + kind=ALTERNATIVE, + description=( + "The household pays its own state income tax before refundable " + "credits during the year, with its Maryland county liability in " + "place of the county estimate, iterated to a fixed point." + ), + plan=_withholding_fixed_point(net=False), + ), + Reading( + id="net", + kind=ALTERNATIVE, + description=( + "As liability, net of state refundable credits and floored at 0." + ), + plan=_withholding_fixed_point(net=True), + ), + ), + ), + UnlistedEstimate( + id="local_sales_tax", + engine_inputs=("local_sales_tax",), + entity="tax_unit", + engine_behavior=( + "local_sales_tax is 0.2 times state_sales_tax (the optional state sales " + "tax table amount; for 2026 the reference system holds the 2025 IRS " + "table) outside CT, DC, IN, KY, MA, MD, ME, MI, NJ and RI, an " + "approximation of the locality's tax " + "(variables/gov/local/tax/sales/local_sales_tax.py). Federal SALT and " + "Hawaii's SALT deduction read it." + ), + why_unlisted="No prompt states a locality or a local sales tax rate.", + found_in="reference_audit 2026-10-05 (variants.py no_local_sales)", + readings=( + Reading( + id="zero", + kind=LITERAL, + description="No local general sales tax (unlisted numeric input = 0).", + plan=_zero_on_tax_unit("local_sales_tax"), + ), + ), + ), + UnlistedEstimate( + id="medicare_part_b_premium", + engine_inputs=("medicare_part_b_premium", "takes_up_medicare_if_eligible"), + entity="person", + engine_behavior=( + "takes_up_medicare_if_eligible defaults to true, so every " + "Medicare-eligible person is medicare_enrolled and is charged a modeled " + "medicare_part_b_premium, net of Medicare Savings Program coverage. " + "medical_expense_health_insurance_premiums adds it unless a direct " + "health_insurance_premiums input is set, which no frozen household does." + ), + why_unlisted=( + "No prompt states Medicare enrollment or a Part B premium, and the " + "prompt says not to infer unlisted expenses or health coverage." + ), + found_in="reference_audit 2026-10-05 (variants.py no_part_b)", + readings=( + Reading( + id="zero", + kind=LITERAL, + description="No Part B premium is paid (unlisted expense = 0).", + plan=_zero_on_people("medicare_part_b_premium"), + ), + Reading( + id="not_enrolled", + kind=ALTERNATIVE, + description=( + "A Medicare-eligible person is not enrolled, so nothing that " + "reads medicare_enrolled applies. The prompt lists no health " + "coverage but also says to assume program take-up, so this is " + "an alternative reading, not the literal one." + ), + plan=_people_flag_alternative( + "takes_up_medicare_if_eligible", False, _is_medicare_eligible + ), + ), + ), + ), + UnlistedEstimate( + id="county", + engine_inputs=("county",), + entity="household", + engine_behavior=( + "A household with no county takes its state's alphabetically first " + "county (first_county_in_state): Allegany County for Maryland, whose " + "county rate then sets the county withholding estimate and liability. " + "County also sets locality flags (in_nyc, in_san_francisco and others) " + "and county-level program rules." + ), + why_unlisted="Prompts state the state, never the county.", + found_in=( + "reference_audit 2026-09-28 (latest_md_local_output_scope) and 2026-10-05" + ), + readings=( + Reading( + id="each_county", + kind=ALTERNATIVE, + description=( + "Each other county of the household's state. A county that " + "switches on a locality flag (living in NYC or San Francisco) " + "makes an unlisted fact true, which the prompt's rule makes " + "false; its moves are marked prompt_rules_out." + ), + plan=_each_county, + ), + ), + ), + UnlistedEstimate( + id="weekly_hours_worked_before_lsr", + engine_inputs=("weekly_hours_worked_before_lsr",), + entity="person", + engine_behavior=( + "SNAP's work rules and the TANF rules of HI, MA, MT, DC and OK read " + "weekly_hours_worked_before_lsr (some through weekly_hours_worked). " + "Stated usual hours reach it through PE_INPUT_ALIASES; otherwise the " + "2.15.17 default is 0 (the literal reading; it was 40 before upstream " + "#9261)." + ), + why_unlisted="Most prompts list no hours worked.", + found_in="reference_audit 2026-09-22 root cause r14 (weekly hours)", + readings=( + Reading( + id="forty_hours", + kind=ALTERNATIVE, + description=( + "A person with no stated hours works 40 hours a week (the " + "1.755.4 default and the 22c board's reading)." + ), + plan=_stated_hours_alternative(40.0), + ), + ), + ), + UnlistedEstimate( + id="mortgage_origination_year", + engine_inputs=( + "first_home_mortgage_origination_year", + "second_home_mortgage_origination_year", + ), + entity="tax_unit", + engine_behavior=( + "An origination year of 0 (the default) takes the post-TCJA " + "acquisition-debt cap ($750,000; $375,000 married filing separately) " + "for the mortgage interest deduction." + ), + why_unlisted=( + "Prompts list mortgage balances and interest, never when the mortgage " + "was taken out (26 U.S.C. 163(h)(3)(F))." + ), + found_in="reference_audit 2026-10-05 follow-up", + readings=( + Reading( + id="pre_tcja", + kind=ALTERNATIVE, + description=( + "The mortgage originated in or before the engine's pre-TCJA " + "origination year (2017), so the $1,000,000 cap applies." + ), + plan=_pre_tcja_mortgage, + ), + ), + ), + UnlistedEstimate( + id="months_receiving_social_security_disability", + engine_inputs=("months_receiving_social_security_disability",), + entity="person", + engine_behavior=( + "The input defaults to 0, so Medicare's under-65 route (24 months of " + "SSDI entitlement) is closed (the literal reading)." + ), + why_unlisted="Prompts list SSDI income, never how long it has been received.", + found_in=( + "reference_exclusions.json (head_medicare_eligible records, 2026-09-05)" + ), + readings=( + Reading( + id="twenty_four_months", + kind=ALTERNATIVE, + description=( + "A person with listed SSDI income has received it for 24 months." + ), + plan=_people_flag_alternative( + "months_receiving_social_security_disability", 24, _receives_ssdi + ), + ), + ), + ), + UnlistedEstimate( + id="meets_ssi_disability_criteria", + engine_inputs=("meets_ssi_disability_criteria",), + entity="person", + engine_behavior=( + "The input defaults to false (the literal reading), so a person listed " + "as disabled does not meet SSI's disability criterion." + ), + why_unlisted=( + "Prompts say a person is disabled, never whether the disability meets " + "the Social Security definition." + ), + found_in="reference_exclusions.json (snap and ssi records, 2026-09-05)", + readings=( + Reading( + id="disabled_meets_criteria", + kind=ALTERNATIVE, + description="A person listed as disabled meets SSI's criterion.", + plan=_people_flag_alternative( + "meets_ssi_disability_criteria", True, _is_disabled + ), + ), + ), + ), +) + +ESTIMATES_BY_ID = {estimate.id: estimate for estimate in ESTIMATES} + + +def select_estimates(ids: Sequence[str] | None) -> tuple[UnlistedEstimate, ...]: + if not ids: + return ESTIMATES + unknown = [i for i in ids if i not in ESTIMATES_BY_ID] + if unknown: + raise SystemExit( + f"Unknown estimate(s): {', '.join(unknown)}. " + f"Registered: {', '.join(ESTIMATES_BY_ID)}" + ) + return tuple(ESTIMATES_BY_ID[i] for i in ids) + + +def combined_literal_plan( + ctx: HouseholdContext, estimates: Sequence[UnlistedEstimate] +) -> list[ReadingPlan]: + """Every literal reading at once: the household exactly as the prompt lists it. + + Some estimates only matter together; the federal SALT deduction falls back to + the local sales tax estimate only once the withholding estimate is gone. + """ + override = Override() + parts = [] + for estimate in estimates: + for reading in estimate.readings: + if reading.kind != LITERAL: + continue + plans = reading.plan(ctx) + if len(plans) > 1 or any(p.variant or p.simulation for p in plans): + raise ValueError( + f"literal reading {estimate.id}/{reading.id} must be one override" + ) + for plan in plans: + if ctx.is_noop(plan.override): + continue + override = override.merged(plan.override) + parts.append(f"{estimate.id}/{reading.id}") + if len(parts) < 2: + return [] + return [ReadingPlan(override, detail={"parts": parts})] + + +# --------------------------------------------------------------------------- +# One household +# --------------------------------------------------------------------------- + + +def build_reference_situation(scenario, engine: Engine, patch=None) -> dict: + """The situation the reference builder simulates for ``scenario``.""" + situation = scenario.to_pe_household() + for person in situation["people"].values(): + for old, new in REFERENCE_INPUT_RENAMES.items(): + if ( + old in person + and not engine.has_variable(old) + and engine.has_variable(new) + ): + person[new] = person.pop(old) + if patch is not None: + situation = patch(copy.deepcopy(situation), scenario) + return situation + + +@dataclass(frozen=True) +class HouseholdJob: + scenario_json: str + variables: tuple[str, ...] + estimate_ids: tuple[str, ...] + combined: bool = True + tolerance: float = TOLERANCE + + +def _reading_record( + estimate_id: str, + reading_id: str, + kind: str, + plan: ReadingPlan, + noop: bool, + outputs: dict[str, float] | None, + localities: Sequence[str] = (), + locality_outputs: Sequence[str] = (), +) -> dict[str, Any]: + return { + "estimate": estimate_id, + "reading": reading_id, + "kind": kind, + "variant": plan.variant, + "override": plan.override.as_dict(), + "override_text": plan.override.describe(), + "noop": noop, + "converged": plan.converged, + "iterations": len(plan.trace), + "trace": plan.trace, + "detail": plan.detail, + "localities": list(localities), + "locality_outputs": list(locality_outputs), + "outputs": outputs, + } + + +def sweep_household(job: HouseholdJob, engine: Engine, patch=None) -> dict[str, Any]: + """The reference and every applicable reading for one household.""" + from policybench.scenarios import scenario_from_dict + + started = time.perf_counter() + scenario = scenario_from_dict(json.loads(job.scenario_json)) + situation = build_reference_situation(scenario, engine, patch) + ctx = HouseholdContext( + scenario=scenario, + situation=situation, + variables=list(job.variables), + engine=engine, + ) + ctx.baseline = engine.simulate(copy.deepcopy(situation)) + baseline_outputs = ctx.outputs(ctx.baseline) + scope_components = dict.fromkeys( + engine.removed_local_components + engine.remaining_local_components + ) + local_taxes = { + name: ctx.total(ctx.baseline, name) + for name in scope_components + if engine.has_variable(name) + } + + def localities_on(simulation: Any) -> set[str]: + return { + flag for flag in engine.locality_flags if ctx.total(simulation, flag) > 0 + } + + baseline_localities = localities_on(ctx.baseline) + + def moved(variable: str, value: float) -> bool: + return output_moved( + baseline_outputs[variable], + value, + binary=is_binary_output(variable), + tolerance=job.tolerance, + ) + + def locality_only( + override: Override, switched_on: Sequence[str], values: dict[str, float] + ) -> list[str]: + """Outputs whose move is the switched-on locality's doing alone. + + The reading is simulated again with every flag it switched on held at its + reference value; an output that then stays put moved only through the + locality's own rules. One that still moves (NY SNAP's utility allowance + region reads the county directly) depends on the county itself. + """ + if any(engine.variable_entity(f) != "household" for f in switched_on): + return [] + held = override.merged(Override(household={f: False for f in switched_on})) + held_values = ctx.outputs(ctx.simulate(held, cache=False)) + return sorted( + variable + for variable, value in values.items() + if moved(variable, value) and not moved(variable, held_values[variable]) + ) + + estimates = [ESTIMATES_BY_ID[i] for i in job.estimate_ids] + readings = [] + output_cache: dict[str, tuple[dict[str, float], list[str], list[str]]] = {} + + def record_reading(estimate_id: str, reading_id: str, kind: str, plan: ReadingPlan): + noop = plan.override.is_empty() or ctx.is_noop(plan.override) + outputs = None + localities: list[str] = [] + locality_outputs: list[str] = [] + if not noop: + key = plan.override.key() + if key not in output_cache: + simulation = plan.simulation or ctx.simulate( + plan.override, cache=not plan.variant + ) + values = ctx.outputs(simulation) + switched_on = sorted(localities_on(simulation) - baseline_localities) + only = locality_only(plan.override, switched_on, values) + output_cache[key] = (values, switched_on, only if switched_on else []) + outputs, localities, locality_outputs = output_cache[key] + readings.append( + _reading_record( + estimate_id, + reading_id, + kind, + plan, + noop, + outputs, + localities, + locality_outputs, + ) + ) + + for estimate in estimates: + for reading in estimate.readings: + for plan in reading.plan(ctx): + record_reading(estimate.id, reading.id, reading.kind, plan) + if job.combined: + for plan in combined_literal_plan(ctx, estimates): + record_reading(COMBINED_ESTIMATE, COMBINED_READING, LITERAL, plan) + + return { + "scenario_id": scenario.id, + "state": scenario.state, + "baseline": baseline_outputs, + "local_taxes": local_taxes, + "baseline_localities": sorted(baseline_localities), + "readings": readings, + "simulations": 1 + ctx.simulations, + "aggregate_local_components": list(engine.aggregate_local_components), + "removed_local_components": list(engine.removed_local_components), + "remaining_local_components": list(engine.remaining_local_components), + "seconds": round(time.perf_counter() - started, 3), + } + + +# --------------------------------------------------------------------------- +# Worker processes +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class EngineConfig: + fix_path: str | None + output_scope_adapter: bool = True + county_states: tuple[str, ...] | None = None + + +def load_fix_module(path: str | Path): + """Import a fix module (``reform`` and/or ``patch(situation, scenario)``).""" + path = Path(path) + digest = hashlib.sha1(str(path).encode()).hexdigest()[:8] + name = f"policybench_fix_{path.stem}_{digest}" + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def compose_reform(fix_reform, adapter: bool): + """The fix module's reform, then the output-scope adapter. + + Returns the composed Reform subclass and a dict that receives, when the system + is built, the local taxes the fix leaves in the state aggregate (the ones the + adapter removes). + """ + from policyengine_core.reforms import Reform + + seen: dict[str, list[str]] = {} + + class composed(Reform): + def apply(self): + if fix_reform is not None: + fix_reform.apply(self) + # The engine can apply a reform more than once; a later pass sees the + # lists the adapter already scoped, so keep the first observation. + seen.setdefault( + "listed", output_scope.all_listed_local_components(self.parameters) + ) + if adapter: + self.modify_parameters(output_scope.remove_local_components) + + return composed, seen + + +def build_engine(config: EngineConfig) -> tuple[Engine, Any]: + from policyengine_us import CountryTaxBenefitSystem, Simulation + + fix_reform = patch = None + if config.fix_path: + module = load_fix_module(config.fix_path) + fix_reform = getattr(module, "reform", None) + patch = getattr(module, "patch", None) + reform, seen = compose_reform(fix_reform, config.output_scope_adapter) + system = CountryTaxBenefitSystem(reform=reform) + listed = list(seen.get("listed", ())) + aggregate = output_scope.listed_local_components(system.parameters) + remaining = output_scope.all_listed_local_components(system.parameters) + removed = [name for name in listed if name not in remaining] + engine = Engine( + system=system, + simulation_class=Simulation, + aggregate_local_components=aggregate, + removed_local_components=removed, + remaining_local_components=remaining, + county_states=( + frozenset(config.county_states) + if config.county_states is not None + else None + ), + ) + return engine, patch + + +_WORKER: dict[str, Any] = {} + + +def _init_worker(config: EngineConfig) -> None: + engine, patch = build_engine(config) + _WORKER["engine"] = engine + _WORKER["patch"] = patch + + +def _run_job(job: HouseholdJob) -> dict[str, Any]: + return sweep_household(job, _WORKER["engine"], _WORKER["patch"]) + + +def _print_progress(done: int, total: int, result: dict[str, Any]) -> None: + print( + f"[{done}/{total}] {result['scenario_id']} {result['state']}: " + f"{result['simulations']} simulations in {result['seconds']:.0f} s", + file=sys.stderr, + flush=True, + ) + + +def run_jobs( + jobs: Sequence[HouseholdJob], + config: EngineConfig, + workers: int, + progress: Callable[[int, int, dict[str, Any]], None] | None = _print_progress, +) -> list[dict[str, Any]]: + """Run every job, building the system once per worker process. + + Results come back in completion order; ``evaluate`` sorts them. + """ + results = [] + + def collect(result: dict[str, Any]) -> None: + results.append(result) + if progress is not None: + progress(len(results), len(jobs), result) + + if workers <= 1: + _init_worker(config) + for job in jobs: + collect(_run_job(job)) + return results + os.environ.setdefault("OPENBLAS_NUM_THREADS", "1") + context = multiprocessing.get_context("spawn") + with ProcessPoolExecutor( + max_workers=workers, + mp_context=context, + initializer=_init_worker, + initargs=(config,), + ) as pool: + futures = [pool.submit(_run_job, job) for job in jobs] + for future in as_completed(futures): + collect(future.result()) + return results + + +# --------------------------------------------------------------------------- +# Comparison and report +# --------------------------------------------------------------------------- + + +def is_binary_output(variable: str) -> bool: + from policybench.spec import metric_type_for_output + + return metric_type_for_output(variable) == "binary" + + +def output_moved( + baseline: float, value: float, *, binary: bool, tolerance: float = TOLERANCE +) -> bool: + """A flag moves when it flips; an amount when it moves past the tolerance.""" + if binary: + return round(baseline) != round(value) + return abs(value - baseline) > tolerance + + +def reproduces(reference: float, value: float, *, binary: bool) -> bool: + if binary: + return round(reference) == round(value) + return abs(value - reference) <= BASELINE_TOLERANCE + + +def exclusion_names_inputs(record: dict | None, inputs: Iterable[str]) -> bool: + """Whether an exclusion record's ``unlisted_input`` names one of ``inputs``. + + Matched as whole identifiers, so ``county`` does not match ``county_fips``. The + field is prose in some records, so a word such as ``county`` still matches + prose that uses it. + """ + if not record: + return False + tokens = set(re.findall(r"[A-Za-z0-9_]+", str(record.get("unlisted_input") or ""))) + return any(name in tokens for name in inputs) + + +def _acknowledgement_key(entry: dict) -> tuple[str, str, str, str]: + return ( + entry["scenario_id"], + entry["variable"], + entry["estimate"], + entry.get("reading", "*"), + ) + + +def load_acknowledgements(path: str | Path | None) -> dict[tuple, dict]: + """Scored moves already under review (a pending decision or ruling), by key. + + The file is ``{"acknowledged": [{"scenario_id", "variable", "estimate", + "reading" (optional; any reading if absent), "status", "note"}]}``. An + ``estimate`` of ``"*"`` covers every reading of the output. + """ + if not path: + return {} + entries = json.loads(Path(path).read_text(encoding="utf-8"))["acknowledged"] + out = {} + for entry in entries: + missing = {"scenario_id", "variable", "estimate", "status"} - set(entry) + if missing: + raise SystemExit( + f"acknowledgement {entry} is missing {', '.join(sorted(missing))}" + ) + out[_acknowledgement_key(entry)] = entry + return out + + +def _acknowledgement_for( + acknowledgements: dict[tuple, dict], + sid: str, + variable: str, + estimate: str, + reading: str, +) -> dict | None: + for key in ( + (sid, variable, estimate, reading), + (sid, variable, estimate, "*"), + (sid, variable, "*", "*"), + ): + if key in acknowledgements: + return acknowledgements[key] + return None + + +def _estimate_inputs(estimate_id: str, parts: Sequence[str] = ()) -> tuple[str, ...]: + if estimate_id == COMBINED_ESTIMATE: + names = [] + for part in parts: + names.extend(ESTIMATES_BY_ID[part.split("/")[0]].engine_inputs) + return tuple(names) + return ESTIMATES_BY_ID[estimate_id].engine_inputs + + +@dataclass +class SweepReport: + baseline: pd.DataFrame + moves: pd.DataFrame + readings: pd.DataFrame + values: pd.DataFrame + summary: dict[str, Any] + + +MOVE_COLUMNS = [ + "scenario_id", + "state", + "variable", + "estimate", + "reading", + "kind", + "variant", + "reference", + "baseline", + "value", + "delta", + "status", + "scored", + "excluded_reason", + "exclusion_names_input", + "acknowledged_status", + "localities", + "override", +] + + +def evaluate( + results: Sequence[dict[str, Any]], + reference: pd.Series, + exclusions: dict[tuple[str, str], dict], + acknowledgements: dict[tuple, dict] | None = None, + tolerance: float = TOLERANCE, +) -> SweepReport: + """Compare every reading with the reference as recomputed, and that with the CSV. + + ``reference`` is the published values indexed by (scenario_id, variable). A move + is measured against the recomputed reference (the swept system's own baseline), + so it shows what the reading alone does. + """ + import pandas as pd + + acknowledgements = acknowledgements or {} + baseline_rows = [] + move_rows = [] + reading_rows = [] + value_rows = [] + for result in sorted(results, key=lambda r: r["scenario_id"]): + sid = result["scenario_id"] + for variable, value in sorted(result["baseline"].items()): + key = (sid, variable) + record = exclusions.get(key) + published = float(reference[key]) + binary = is_binary_output(variable) + baseline_rows.append( + { + "scenario_id": sid, + "state": result["state"], + "variable": variable, + "reference": published, + "baseline": value, + "reproduces": reproduces(published, value, binary=binary), + "scored": record is None, + "excluded_reason": record["reason_code"] if record else "", + } + ) + for reading in result["readings"]: + moved_here = 0 + parts = reading["detail"].get("parts", ()) + inputs = _estimate_inputs(reading["estimate"], parts) + outputs = reading["outputs"] + if outputs is not None: + for variable, value in sorted(outputs.items()): + base = result["baseline"][variable] + value_rows.append( + { + "scenario_id": sid, + "variable": variable, + "estimate": reading["estimate"], + "reading": reading["reading"], + "variant": reading["variant"], + "baseline": base, + "value": value, + } + ) + binary = is_binary_output(variable) + if not output_moved( + base, value, binary=binary, tolerance=tolerance + ): + continue + moved_here += 1 + key = (sid, variable) + record = exclusions.get(key) + names_input = exclusion_names_inputs(record, inputs) + ack = _acknowledgement_for( + acknowledgements, + sid, + variable, + reading["estimate"], + reading["reading"], + ) + if record is not None: + status = ( + "excluded_same_input" + if names_input + else "excluded_other_reason" + ) + elif variable in reading.get("locality_outputs", ()): + status = "prompt_rules_out" + elif ack is not None: + status = "acknowledged" + else: + status = "scored" + move_rows.append( + { + "scenario_id": sid, + "state": result["state"], + "variable": variable, + "estimate": reading["estimate"], + "reading": reading["reading"], + "kind": reading["kind"], + "variant": reading["variant"], + "reference": float(reference[key]), + "baseline": base, + "value": value, + "delta": value - base, + "status": status, + "scored": record is None, + "excluded_reason": record["reason_code"] if record else "", + "exclusion_names_input": names_input, + "acknowledged_status": ack["status"] if ack else "", + "localities": ";".join(reading["localities"]), + "override": reading["override_text"], + } + ) + reading_rows.append( + { + "scenario_id": sid, + "state": result["state"], + "estimate": reading["estimate"], + "reading": reading["reading"], + "kind": reading["kind"], + "variant": reading["variant"], + "noop": reading["noop"], + "converged": reading["converged"], + "iterations": reading["iterations"], + "outputs_moved": moved_here, + "localities": ";".join(reading["localities"]), + "override": reading["override_text"], + "detail": json.dumps(reading["detail"], sort_keys=True), + } + ) + + baseline = pd.DataFrame(baseline_rows) + moves = pd.DataFrame(move_rows, columns=MOVE_COLUMNS) + readings = pd.DataFrame(reading_rows) + values = pd.DataFrame(value_rows) + summary = summarize(results, baseline, moves, readings) + return SweepReport(baseline, moves, readings, values, summary) + + +def summarize( + results: Sequence[dict[str, Any]], + baseline: pd.DataFrame, + moves: pd.DataFrame, + readings: pd.DataFrame, +) -> dict[str, Any]: + + mismatched = baseline[~baseline["reproduces"] & baseline["scored"]] + scope = [ + {"scenario_id": r["scenario_id"], "local_taxes": r["local_taxes"]} + for r in results + if any(abs(v) > 0.005 for v in r["local_taxes"].values()) + ] + unconverged = readings[~readings["converged"]] if not readings.empty else readings + per_reading = [] + if not readings.empty: + grouped = readings.groupby(["estimate", "reading", "kind"], sort=True) + for (estimate, reading, kind), frame in grouped: + active = frame[~frame["noop"]] + these = moves[ + (moves["estimate"] == estimate) & (moves["reading"] == reading) + ] + per_reading.append( + { + "estimate": estimate, + "reading": reading, + "kind": kind, + "households_planned": int(frame["scenario_id"].nunique()), + "households_changed": int(active["scenario_id"].nunique()), + "simulated_variants": int(len(active)), + "outputs_moved": int( + these[["scenario_id", "variable"]].drop_duplicates().shape[0] + ), + "scored_outputs_moved": int( + these[these["status"] == "scored"][["scenario_id", "variable"]] + .drop_duplicates() + .shape[0] + ), + } + ) + by_output = [] + if not moves.empty: + for (sid, variable), frame in moves.groupby( + ["scenario_id", "variable"], sort=True + ): + by_output.append( + { + "scenario_id": sid, + "variable": variable, + "status": _strongest_status(frame["status"]), + "excluded_reason": frame["excluded_reason"].iloc[0], + "estimates": sorted(set(frame["estimate"])), + "localities": sorted( + { + flag + for text in frame["localities"].fillna("") + for flag in str(text).split(";") + if flag + } + ), + "readings": sorted(set(frame["estimate"] + "/" + frame["reading"])), + "reference": float(frame["reference"].iloc[0]), + "baseline": float(frame["baseline"].iloc[0]), + "min_value": float(frame["value"].min()), + "max_value": float(frame["value"].max()), + } + ) + return { + "outputs": int(len(baseline)), + "scored_outputs": int(baseline["scored"].sum()) if len(baseline) else 0, + "households": int(len(results)), + "simulations": int(sum(r["simulations"] for r in results)), + "baseline_scored_mismatches": mismatched[ + ["scenario_id", "variable", "reference", "baseline"] + ].to_dict("records"), + "baseline_excluded_mismatches": int( + (~baseline["reproduces"] & ~baseline["scored"]).sum() + ) + if len(baseline) + else 0, + "scope_violations": scope, + "unconverged": unconverged[["scenario_id", "estimate", "reading"]].to_dict( + "records" + ) + if len(unconverged) + else [], + "per_reading": per_reading, + "moved_outputs": by_output, + "counts": { + status: sum(1 for o in by_output if o["status"] == status) + for status in _STATUS_ORDER + }, + } + + +# Strongest first: an output's status is the strongest of its moves. A move is +# scored unless the output is excluded, the reading switches on a locality fact the +# prompt's rule makes false (prompt_rules_out), or an --acknowledged entry covers it. +_STATUS_ORDER = ( + "scored", + "acknowledged", + "prompt_rules_out", + "excluded_other_reason", + "excluded_same_input", +) + + +def _strongest_status(statuses: Iterable[str]) -> str: + present = set(statuses) + for status in _STATUS_ORDER: + if status in present: + return status + return "" + + +def render_markdown(summary: dict[str, Any], moves: pd.DataFrame) -> str: + meta = summary["run"] + lines = [ + "# Unlisted-input sweep", + "", + f"- Engine: policyengine-us {meta['policyengine_us']} " + f"(policyengine-core {meta['policyengine_core']})", + f"- Reference system: `{meta['fix_id']}`" + + (" + `policybench_output_scope`" if meta["output_scope_adapter"] else ""), + f"- Run: {meta['scenarios']} ({summary['households']} households, " + f"{summary['outputs']} outputs, {summary['scored_outputs']} scored)", + f"- Simulations: {summary['simulations']} in {meta['seconds']:.0f} s " + f"on {meta['workers']} workers", + f"- Tolerance: ${meta['tolerance']:g} for amounts; flags move when they flip", + "", + "## Gates", + "", + f"- Baseline reproduces every scored reference: " + f"{'yes' if not summary['baseline_scored_mismatches'] else 'NO'} " + f"({len(summary['baseline_scored_mismatches'])} mismatches)", + f"- Local taxes and credits in the state outputs: " + f"{'none' if not summary['scope_violations'] else 'FOUND'}" + f" (removed by the adapter and zero in every household: " + f"{', '.join(meta['removed_local_components']) or 'none'}; " + f"still listed: {', '.join(meta['remaining_local_components']) or 'none'})", + f"- Every reference row recomputed: " + f"{'yes' if not summary.get('reference_rows_not_recomputed') else 'NO'}", + f"- Fixed points converged: {'yes' if not summary['unconverged'] else 'NO'}", + f"- Scored outputs that move: {summary['counts']['scored']}", + f"- Outputs that move only where a reading puts the household in a " + f"locality (prompt_rules_out): {summary['counts']['prompt_rules_out']}", + "", + "## Readings", + "", + "| Estimate | Reading | Kind | Households changed | Simulations | " + "Outputs moved | Scored |", + "|---|---|---|---:|---:|---:|---:|", + ] + for row in summary["per_reading"]: + lines.append( + f"| {row['estimate']} | {row['reading']} | {row['kind']} | " + f"{row['households_changed']} | {row['simulated_variants']} | " + f"{row['outputs_moved']} | {row['scored_outputs_moved']} |" + ) + lines += ["", "## Outputs that move", ""] + if not summary["moved_outputs"]: + lines.append("None.") + else: + lines += [ + "| Output | Status | Published | Recomputed | Range under readings " + "| Readings |", + "|---|---|---:|---:|---|---|", + ] + for row in summary["moved_outputs"]: + status = row["status"] + if row["excluded_reason"]: + status += f" ({row['excluded_reason']})" + elif row["status"] == "prompt_rules_out": + status += f" ({', '.join(row['localities'])})" + lines.append( + f"| {row['scenario_id']} {row['variable']} | {status} | " + f"{row['reference']:,.2f} | {row['baseline']:,.2f} | " + f"{row['min_value']:,.2f} to " + f"{row['max_value']:,.2f} | {', '.join(row['readings'])} |" + ) + lines.append("") + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Command +# --------------------------------------------------------------------------- + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _package_version(name: str) -> str: + try: + return version(name) + except PackageNotFoundError: + return "not installed" + + +def assemble_fix_dir(fix: Path, support: Sequence[Path], destination: Path) -> Path: + """Copy the fix module's directory and support files into ``destination``. + + Fix modules load their parts and data files from their own directory + (``Path(__file__).with_name``); the 2026-09-28 sales tax convention reads + ``r19_irs_sales_tax_2025.json``, which lives in the 2026-09-22 directory. + """ + destination.mkdir(parents=True, exist_ok=True) + for path in sorted(fix.parent.iterdir()): + if path.is_file() and path.suffix in {".py", ".json"}: + shutil.copy2(path, destination / path.name) + for path in support: + shutil.copy2(path, destination / path.name) + return destination / fix.name + + +def _county_count(state: str) -> int: + try: + from policyengine_us.variables.household.demographic.geographic.county.county_enum import ( # noqa: E501 + County, + ) + except ImportError: # pragma: no cover - the engine moved the enum + return 0 + return sum(1 for c in County if str(c.value).endswith(f", {state}")) + + +def build_jobs( + scenarios: pd.DataFrame, + reference: pd.Series, + programs: Sequence[str], + estimates: Sequence[UnlistedEstimate], + county_states: frozenset[str] | None, + combined: bool = True, + tolerance: float = TOLERANCE, +) -> list[HouseholdJob]: + """One job per household, longest first so the pool's tail is short.""" + from policybench.scenarios import scenario_from_dict + from policybench.spec import expand_programs_for_scenario + + estimate_ids = tuple(e.id for e in estimates) + jobs = [] + for _, row in scenarios.iterrows(): + scenario = scenario_from_dict(json.loads(row["scenario_json"])) + variables = tuple( + v + for v in expand_programs_for_scenario(list(programs), scenario) + if (scenario.id, v) in reference.index + ) + cost = 1 + if "county" in estimate_ids and ( + county_states is None or scenario.state in county_states + ): + cost += _county_count(scenario.state) + jobs.append( + ( + cost, + scenario.id, + HouseholdJob( + row["scenario_json"], variables, estimate_ids, combined, tolerance + ), + ) + ) + jobs.sort(key=lambda item: (-item[0], item[1])) + return [job for _, _, job in jobs] + + +def _resolve(run_dir: Path, explicit: str | None, name: str) -> Path: + return Path(explicit) if explicit else run_dir / name + + +def run(args) -> int: + """The ``unlisted-input-sweep`` command. Returns the process exit code.""" + import pandas as pd + + started = time.perf_counter() + run_dir = Path(args.run_dir) + scenarios_path = _resolve(run_dir, args.scenarios, "scenarios.csv") + reference_path = _resolve(run_dir, args.reference, "reference_outputs.csv") + meta_path = Path(f"{reference_path}.meta.json") + exclusions_path = _resolve(run_dir, args.exclusions, "reference_exclusions.json") + for path in (scenarios_path, reference_path, meta_path): + if not path.exists(): + raise SystemExit(f"Missing {path}") + + programs = json.loads(meta_path.read_text(encoding="utf-8"))["programs"] + reference = pd.read_csv(reference_path).set_index(["scenario_id", "variable"])[ + "value" + ] + exclusions = {} + if exclusions_path.exists(): + for record in json.loads(exclusions_path.read_text(encoding="utf-8"))[ + "exclusions" + ]: + exclusions[(record["scenario_id"], record["variable"])] = record + acknowledgements = load_acknowledgements(args.acknowledged) + scenarios = pd.read_csv(scenarios_path) + if args.scenario: + scenarios = scenarios[scenarios["scenario_id"].isin(args.scenario)] + if scenarios.empty: + raise SystemExit("No scenarios match --scenario") + if (scenarios["country"] != "us").any(): + raise SystemExit("The unlisted-input sweep covers US scenarios only") + scope_year = int(output_scope.SCOPE_INSTANT[:4]) + years = { + int(json.loads(raw).get("year", scope_year)) + for raw in scenarios["scenario_json"] + } + if years != {scope_year}: + raise SystemExit( + f"Scenario years {sorted(years)}: the output-scope adapter rewrites the " + f"{scope_year} lists only" + ) + + estimates = select_estimates(args.estimate) + county_states = ( + None + if args.county_states in (None, "all") + else frozenset(s.strip().upper() for s in args.county_states.split(",")) + ) + jobs = build_jobs( + scenarios, + reference, + programs, + estimates, + county_states, + combined=not args.no_combined, + tolerance=args.tolerance, + ) + # Every reference row of the swept households must be recomputed, or the + # baseline gate would pass over the rows a program expansion dropped. + swept = set(scenarios["scenario_id"]) + expected = {key for key in reference.index if key[0] in swept} + planned = { + (json.loads(job.scenario_json)["id"], variable) + for job in jobs + for variable in job.variables + } + not_recomputed = sorted(expected - planned) + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + fix = Path(args.fix) if args.fix else None + support = [Path(p) for p in args.fix_support or []] + workers = args.workers or max(1, min(16, (os.cpu_count() or 2) - 2)) + with tempfile.TemporaryDirectory(prefix="unlisted_input_sweep_fix_") as tmp: + fix_files = {} + fix_path = None + if fix is not None: + fix_path = assemble_fix_dir(fix, support, Path(tmp)) + fix_files = { + p.name: _sha256(p) for p in sorted(Path(tmp).iterdir()) if p.is_file() + } + config = EngineConfig( + fix_path=str(fix_path) if fix_path else None, + output_scope_adapter=not args.no_output_scope_adapter, + county_states=tuple(sorted(county_states)) if county_states else None, + ) + print( + f"policyengine-us {_package_version('policyengine-us')}; " + f"{len(jobs)} households; {len(estimates)} estimates; {workers} workers", + flush=True, + ) + results = run_jobs(jobs, config, workers) + scope = scope_lists(results) + + report = evaluate(results, reference, exclusions, acknowledgements, args.tolerance) + summary = report.summary + summary["reference_rows_not_recomputed"] = [list(key) for key in not_recomputed] + summary["run"] = { + "generated_at_utc": datetime.now(timezone.utc).isoformat(), + "policyengine_us": _package_version("policyengine-us"), + "policyengine_core": _package_version("policyengine-core"), + "fix": str(fix) if fix else None, + "fix_id": fix.stem if fix else "baseline", + "fix_files_sha256": fix_files, + "output_scope_adapter": not args.no_output_scope_adapter, + **scope, + "scenarios": str(scenarios_path), + "scenarios_sha256": _sha256(scenarios_path), + "reference": str(reference_path), + "reference_sha256": _sha256(reference_path), + "exclusions": str(exclusions_path) if exclusions_path.exists() else None, + "exclusions_sha256": _sha256(exclusions_path) + if exclusions_path.exists() + else None, + "acknowledged": args.acknowledged, + "tolerance": args.tolerance, + "county_states": sorted(county_states) if county_states else "all", + "combined_literal_reading": not args.no_combined, + "workers": workers, + "seconds": round(time.perf_counter() - started, 1), + } + summary["estimates"] = [e.describe() for e in estimates] + + report.moves.to_csv(out_dir / "moves.csv", index=False) + report.readings.to_csv(out_dir / "readings.csv", index=False) + report.baseline.to_csv(out_dir / "baseline.csv", index=False) + report.values.to_csv(out_dir / "values.csv.gz", index=False) + (out_dir / "summary.json").write_text( + json.dumps(summary, indent=1, sort_keys=True, default=_json_default) + "\n", + encoding="utf-8", + ) + (out_dir / "households.json").write_text( + json.dumps( + [ + {k: v for k, v in r.items() if k != "readings"} + | { + "readings": [ + {k: v for k, v in reading.items() if k != "outputs"} + for reading in r["readings"] + if not reading["noop"] + ] + } + for r in sorted(results, key=lambda r: r["scenario_id"]) + ], + indent=1, + sort_keys=True, + default=_json_default, + ) + + "\n", + encoding="utf-8", + ) + markdown = render_markdown(summary, report.moves) + (out_dir / "report.md").write_text(markdown, encoding="utf-8") + print(markdown) + return exit_code( + summary, + strict=args.strict, + allow_baseline_mismatch=args.allow_baseline_mismatch, + ) + + +SCOPE_KEYS = ( + "aggregate_local_components", + "removed_local_components", + "remaining_local_components", +) + + +def scope_lists(results: Sequence[dict[str, Any]]) -> dict[str, list[str]]: + """The local entries the swept system removed from, or keeps in, its state lists. + + Every worker builds the same system, so every household reports the same lists. + """ + lists = {tuple(tuple(r[key]) for key in SCOPE_KEYS) for r in results} + if len(lists) > 1: + raise RuntimeError(f"workers built different systems: {sorted(lists)}") + if not lists: + return {key: [] for key in SCOPE_KEYS} + return {key: list(values) for key, values in zip(SCOPE_KEYS, lists.pop())} + + +def exit_code( + summary: dict[str, Any], *, strict: bool, allow_baseline_mismatch: bool +) -> int: + """1 when the sweep itself fails a gate; 2 when --strict finds scored moves.""" + if summary["baseline_scored_mismatches"] and not allow_baseline_mismatch: + return 1 + if summary["scope_violations"] or summary["unconverged"]: + return 1 + if summary.get("reference_rows_not_recomputed"): + return 1 + if strict and summary["counts"]["scored"]: + return 2 + return 0 + + +def _json_default(value: Any) -> Any: + if isinstance(value, float) and math.isnan(value): + return None + if hasattr(value, "item"): + return value.item() + raise TypeError(f"not JSON serializable: {type(value).__name__}") + + +def add_arguments(parser) -> None: + parser.add_argument( + "--fix", + default=None, + help=( + "Fix module defining the reference system (reform and/or " + "patch(situation, scenario)), e.g. " + "reference_audit/2026-09-28/fixes/latest_final.py" + ), + ) + parser.add_argument( + "--fix-support", + action="append", + default=None, + help="Extra file the fix module reads from its directory (repeatable)", + ) + parser.add_argument( + "--run-dir", + default=str(DEFAULT_RUN_DIR), + help="Frozen run directory with scenarios.csv, reference_outputs.csv, " + "its .meta.json and reference_exclusions.json", + ) + parser.add_argument("--scenarios", default=None, help="Override scenarios.csv") + parser.add_argument( + "--reference", default=None, help="Override reference_outputs.csv" + ) + parser.add_argument( + "--exclusions", default=None, help="Override reference_exclusions.json" + ) + parser.add_argument("--out-dir", required=True) + parser.add_argument( + "--estimate", + action="append", + default=None, + help=f"Registered estimate to sweep (repeatable; default all): " + f"{', '.join(ESTIMATES_BY_ID)}", + ) + parser.add_argument( + "--scenario", action="append", default=None, help="Restrict to a scenario id" + ) + parser.add_argument( + "--county-states", + default="all", + help="States whose counties the county estimate sweeps: 'all' (default) " + "or a comma-separated list such as MD", + ) + parser.add_argument("--workers", type=int, default=None) + parser.add_argument("--tolerance", type=float, default=TOLERANCE) + parser.add_argument( + "--acknowledged", + default=None, + help="JSON of scored moves already under review, so --strict passes them", + ) + parser.add_argument( + "--strict", + action="store_true", + help="Exit 2 when any scored output moves without an exclusion or " + "acknowledgement", + ) + parser.add_argument( + "--allow-baseline-mismatch", + action="store_true", + help="Do not fail when the recomputed reference differs from the CSV " + "(previewing a new engine)", + ) + parser.add_argument( + "--no-output-scope-adapter", + action="store_true", + help="Sweep the fix module's system without policybench.output_scope", + ) + parser.add_argument( + "--no-combined", + action="store_true", + help="Skip the reading that applies every literal reading at once", + ) diff --git a/reference_audit/unlisted_input_acknowledged.json b/reference_audit/unlisted_input_acknowledged.json new file mode 100644 index 00000000..5b1e79f8 --- /dev/null +++ b/reference_audit/unlisted_input_acknowledged.json @@ -0,0 +1,33 @@ +{ + "note": "Scored outputs of the current references that the unlisted-input sweep finds moving and that await a ruling. Pass this file to `policybench unlisted-input-sweep --acknowledged`. Remove an entry when the release that rules on it lands (tests/test_unlisted_input_sweep.py fails once its output is excluded).", + "acknowledged": [ + { + "scenario_id": "scenario_022", + "variable": "federal_income_tax_before_refundable_credits", + "estimate": "*", + "status": "pending d963", + "note": "PolicyEngine/policybench#191 proposes excluding it: the federal SALT deduction rests on the engine's withholding estimate (state_withheld_income_tax)." + }, + { + "scenario_id": "scenario_081", + "variable": "federal_income_tax_before_refundable_credits", + "estimate": "*", + "status": "pending d963", + "note": "PolicyEngine/policybench#191 proposes excluding it: the federal SALT deduction rests on the engine's withholding estimate (state_withheld_income_tax)." + }, + { + "scenario_id": "scenario_114", + "variable": "federal_income_tax_before_refundable_credits", + "estimate": "*", + "status": "pending d963", + "note": "PolicyEngine/policybench#191 proposes excluding it: the federal SALT deduction rests on the engine's withholding estimate; the modeled Medicare Part B premium also moves it." + }, + { + "scenario_id": "scenario_114", + "variable": "state_income_tax_before_refundable_credits", + "estimate": "*", + "status": "pending d974", + "note": "PolicyEngine/policybench#193 proposes excluding it: the Virginia medical deduction counts the engine's modeled Medicare Part B premium." + } + ] +} diff --git a/tests/test_output_scope.py b/tests/test_output_scope.py new file mode 100644 index 00000000..d672e11c --- /dev/null +++ b/tests/test_output_scope.py @@ -0,0 +1,195 @@ +"""Tests for policybench/output_scope.py: local taxes stay out of the state output.""" + +from __future__ import annotations + +import numpy as np +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +from policybench import output_scope + +STATE_ENTRIES = [ + "ca_income_tax_before_refundable_credits", + "md_income_tax_before_refundable_credits", + "ny_income_tax_before_refundable_credits", + "mt_income_tax_before_refundable_credits_unit", + "nc_income_tax", +] + + +def test_known_local_taxes_are_recognized(): + for name in output_scope.LOCAL_INCOME_TAX_COMPONENTS: + assert output_scope.is_local_component(name) + for name in STATE_ENTRIES: + assert not output_scope.is_local_component(name) + assert output_scope.is_local_component("pa_philadelphia_wage_tax") + assert output_scope.is_local_component("in_local_income_tax_before_credits") + assert output_scope.is_local_component("ca_sf_wftc") + # SALT variables carry "state_and_local" and are state entries. + assert not output_scope.is_local_component( + "va_capped_state_and_local_sales_or_income_tax" + ) + + +@settings(max_examples=200, deadline=None) +@given( + st.lists( + st.sampled_from(STATE_ENTRIES + list(output_scope.LOCAL_INCOME_TAX_COMPONENTS)), + max_size=12, + ) +) +def test_scoping_keeps_every_state_entry_in_order_and_no_local_one(entries): + kept = output_scope.scoped_components(entries) + assert kept == [e for e in entries if e in STATE_ENTRIES] + assert not output_scope.local_components(kept) + assert output_scope.scoped_components(kept) == kept + assert sorted(kept + output_scope.local_components(entries)) == sorted(entries) + + +class Leaf: + def __init__(self, value): + self.value = list(value) + self.updates = [] + + def __call__(self, instant): + return list(self.value) + + def update(self, period, value): + self.updates.append(period) + self.value = list(value) + + +class Node: + pass + + +def tree(aggregate, refundable=None, ca_refundable=None): + """A parameter tree holding the three scoped lists.""" + root = Node() + leaves = { + output_scope.STATE_AGGREGATE_PARAMETER: Leaf(aggregate), + output_scope.STATE_REFUNDABLE_PARAMETER: Leaf( + refundable or ["ny_refundable_credits", "ca_refundable_credits"] + ), + output_scope.CA_REFUNDABLE_PARAMETER: Leaf(ca_refundable or ["ca_eitc"]), + } + for path, leaf in leaves.items(): + node = root + parts = path.split(".") + for part in parts[:-1]: + if not hasattr(node, part): + setattr(node, part, Node()) + node = getattr(node, part) + setattr(node, parts[-1], leaf) + return root, leaves + + +def test_remove_local_components_clears_every_list_and_updates_only_when_needed(): + params, leaves = tree( + STATE_ENTRIES + ["nyc_income_tax_before_refundable_credits"], + ["ny_refundable_credits", "nyc_refundable_credits", "ca_refundable_credits"], + ["ca_eitc", "ca_sf_wftc"], + ) + assert output_scope.listed_local_components(params) == [ + "nyc_income_tax_before_refundable_credits" + ] + assert output_scope.all_listed_local_components(params) == [ + "nyc_income_tax_before_refundable_credits", + "nyc_refundable_credits", + "ca_sf_wftc", + ] + output_scope.remove_local_components(params) + aggregate = leaves[output_scope.STATE_AGGREGATE_PARAMETER] + assert aggregate.value == STATE_ENTRIES + assert aggregate.updates == [output_scope.SCOPE_PERIOD] + assert leaves[output_scope.STATE_REFUNDABLE_PARAMETER].value == [ + "ny_refundable_credits", + "ca_refundable_credits", + ] + assert leaves[output_scope.CA_REFUNDABLE_PARAMETER].value == ["ca_eitc"] + assert output_scope.all_listed_local_components(params) == [] + output_scope.remove_local_components(params) + assert aggregate.updates == [output_scope.SCOPE_PERIOD] + + +def test_remove_local_components_refuses_to_drop_a_state_entry(): + params, _ = tree(["ca_income_tax_before_refundable_credits"]) + with pytest.raises(AssertionError): + output_scope.remove_local_components(params) + + +class Sim: + def __init__(self, values): + self.values = values + + def calculate(self, name, period): + return np.array([self.values[name]]) + + +def test_scope_violations_lists_each_nonzero_local_tax(): + sim = Sim( + { + "nyc_income_tax_before_refundable_credits": 1_234.5, + "md_local_income_tax_before_refundable_credits": 0.0, + } + ) + assert output_scope.scope_violations( + sim, output_scope.LOCAL_INCOME_TAX_COMPONENTS, 2026 + ) == {"nyc_income_tax_before_refundable_credits": 1_234.5} + + +def test_reform_is_a_policyengine_reform_built_on_first_use(): + from policyengine_core.reforms import Reform + + assert issubclass(output_scope.reform, Reform) + assert output_scope.reform is output_scope.reform + with pytest.raises(AttributeError): + output_scope.no_such_attribute # noqa: B018 + + +@pytest.mark.parametrize("path", list(output_scope.SCOPED_LISTS)) +def test_installed_engine_lists_no_local_entry_the_adapter_misses(path): + """Guards every policyengine-us upgrade: a new local entry in a state list must + be named here, and every state entry the adapter keeps must still be listed.""" + entries = output_scope.installed_parameter_list(path) + spec = output_scope.SCOPED_LISTS[path] + for name in spec["keeps"]: + assert name in entries + local = output_scope.local_components(entries) + assert local == [e for e in entries if e in spec["local"]], local + assert set(spec["local"]) <= set(local) + assert not output_scope.local_components(output_scope.scoped_components(entries)) + + +def test_no_other_state_list_carries_a_local_entry(): + """Every list feeding the two state outputs is scoped: the state income tax and + state refundable credit lists, and each state's refundable list they add.""" + import importlib.util + from pathlib import Path + + import yaml + + spec = importlib.util.find_spec("policyengine_us") + root = Path(next(iter(spec.submodule_search_locations))) / "parameters" + scoped = {path.replace(".", "/") + ".yaml" for path in output_scope.SCOPED_LISTS} + hits = {} + for path in root.glob("gov/states/**/*.yaml"): + relative = str(path.relative_to(root)) + data = yaml.load(path.read_text(encoding="utf-8"), Loader=yaml.BaseLoader) + if not isinstance(data, dict) or not isinstance(data.get("values"), dict): + continue + latest = data["values"][sorted(data["values"])[-1]] + if not isinstance(latest, list): + continue + local = output_scope.local_components(str(x) for x in latest) + # Refundable lists feed state_refundable_credits and each state's + # nonrefundable list reduces its income tax before refundable credits. The + # household-level state_non_refundable_credits aggregate feeds no output. + feeds_output = ( + "refundable" in relative or "income_tax_before_refundable" in relative + ) and not relative.endswith("household/state_non_refundable_credits.yaml") + if local and feeds_output: + hits["gov/" + relative.split("gov/", 1)[-1]] = local + unscoped = {k: v for k, v in hits.items() if k not in scoped} + assert not unscoped, unscoped diff --git a/tests/test_unlisted_input_sweep.py b/tests/test_unlisted_input_sweep.py new file mode 100644 index 00000000..55d7daf4 --- /dev/null +++ b/tests/test_unlisted_input_sweep.py @@ -0,0 +1,1814 @@ +"""Tests for the unlisted-input sweep (policybench/unlisted_input_sweep.py). + +Every test here runs on a small fake engine: FakeSimulation computes a handful of +tax and benefit variables from the situation dict with toy formulas shaped like the +engine's (the SALT deduction takes the larger of income and sales tax, a state's +tax reads the federal SALT deduction, county sets Maryland's local rate). The +real-engine checks are marked slow. +""" + +from __future__ import annotations + +import copy +import json +import re +from enum import Enum +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +from policybench import unlisted_input_sweep as sweep +from policybench.scenarios import ( + Person, + Scenario, + is_excluded_prompt_input_name, + scenario_manifest, + scenario_to_dict, +) + +YEAR = "2026" +INCOME_TAX_STATES = {"MD", "CA"} +NO_LOCAL_SALES_TAX_STATES = {"CT"} +MD_RATES = { + "ALLEGANY_COUNTY_MD": 0.0303, + "BALTIMORE_CITY_MD": 0.032, + "WORCESTER_COUNTY_MD": 0.0225, +} +PROGRAMS = [ + "federal_income_tax_before_refundable_credits", + "state_income_tax_before_refundable_credits", + "snap", + "ssi", + "person_medicare_eligible", +] + + +class FakeCounty(Enum): + UNKNOWN = "UNKNOWN" + WORCESTER_COUNTY_MD = "Worcester County, MD" + ALLEGANY_COUNTY_MD = "Allegany County, MD" + BALTIMORE_CITY_MD = "Baltimore city, MD" + ALPINE_COUNTY_CA = "Alpine County, CA" + ALAMEDA_COUNTY_CA = "Alameda County, CA" + ALBANY_COUNTY_NY = "Albany County, NY" + KINGS_COUNTY_NY = "Kings County, NY" + NASSAU_COUNTY_NY = "Nassau County, NY" + + +FIRST_COUNTY = { + "MD": "ALLEGANY_COUNTY_MD", + "CA": "ALAMEDA_COUNTY_CA", + "NY": "ALBANY_COUNTY_NY", +} +NYC_COUNTIES = {"KINGS_COUNTY_NY"} +MCTD_ZONE_2 = {"NASSAU_COUNTY_NY"} +# A toy SNAP utility allowance region that reads the county directly, as New +# York's does: NYC and Nassau-Suffolk get higher allowances. +SNAP_REGION_BONUS = {"KINGS_COUNTY_NY": 300.0, "NASSAU_COUNTY_NY": 200.0} + + +class EnumResult(list): + def decode_to_str(self): + return np.array(self, dtype=object) + + +class Node(SimpleNamespace): + """A parameter tree node: attributes for children, a value when called.""" + + def __init__(self, value=None, **children): + super().__init__(**children) + self._value = value + self.updates = [] + + def __call__(self, instant): + return self._value + + def update(self, period, value): + self.updates.append((period, value)) + self._value = value + + +def fake_parameters( + aggregate, + refundable=("ny_refundable_credits", "ca_refundable_credits"), + ca_refundable=("ca_eitc",), +): + return Node( + gov=Node( + irs=Node( + deductions=Node( + itemized=Node( + interest=Node( + mortgage=Node(pre_tcja_origination_year=Node(2017)) + ) + ) + ) + ), + states=Node( + household=Node( + state_income_tax_before_refundable_credits=Node(list(aggregate)), + state_refundable_credits=Node(list(refundable)), + ), + ca=Node( + tax=Node( + income=Node(credits=Node(refundable=Node(list(ca_refundable)))) + ) + ), + ), + ) + ) + + +class FakeSystem: + def __init__(self, aggregate=("md_income_tax_before_refundable_credits",)): + self.aggregate = list(aggregate) + self.parameters = fake_parameters(aggregate) + names = [n[3:] for n in dir(FakeSimulation) if n.startswith("_v_")] + self.variables = {name: SimpleNamespace() for name in names} + self.variables["county"] = SimpleNamespace(possible_values=FakeCounty) + self.variables["partnership_self_employment_net_earnings"] = SimpleNamespace() + + +class FakeSimulation: + """Toy formulas over a PolicyBench situation; an input overrides its formula.""" + + instances = 0 + state_salt_share = -0.05 + salt_cap = 10_000.0 + + def __init__(self, tax_benefit_system, situation): + type(self).instances += 1 + self.system = tax_benefit_system + self.situation = situation + self.people = list(situation["people"].values()) + self.tax_unit = next(iter(situation["tax_units"].values())) + self.household = next(iter(situation["households"].values())) + self.cache = {} + + def calculate(self, variable, period): + assert str(period) == YEAR + if variable not in self.cache: + self.cache[variable] = getattr(self, f"_v_{variable}")() + return self.cache[variable] + + # helpers + def _person(self, name, default): + return np.array( + [float(p.get(name, {YEAR: default})[YEAR]) for p in self.people] + ) + + def _unit(self, node, name, formula): + if name in node: + return np.array([float(node[name][YEAR])]) + return np.array([float(formula())]) + + def _sum(self, name): + return float(self.calculate(name, YEAR).sum()) + + @property + def state(self): + return self.household["state_code"][YEAR] + + # person variables + def _v_age(self): + return self._person("age", 0) + + def _v_employment_income(self): + return self._person("employment_income", 0) + + def _v_social_security_disability(self): + return self._person("social_security_disability", 0) + + def _v_is_disabled(self): + return self._person("is_disabled", False).astype(bool) + + def _v_weekly_hours_worked_before_lsr(self): + return self._person("weekly_hours_worked_before_lsr", 0) + + def _v_months_receiving_social_security_disability(self): + return self._person("months_receiving_social_security_disability", 0) + + def _v_meets_ssi_disability_criteria(self): + return self._person("meets_ssi_disability_criteria", False).astype(bool) + + def _v_medicare_part_b_premium(self): + enrolled = (self.calculate("age", YEAR) >= 65) & self.calculate( + "takes_up_medicare_if_eligible", YEAR + ) + modeled = np.where(enrolled, 2_000.0, 0.0) + return np.array( + [ + float(p["medicare_part_b_premium"][YEAR]) + if "medicare_part_b_premium" in p + else modeled[i] + for i, p in enumerate(self.people) + ] + ) + + def _v_is_medicare_eligible(self): + age = self.calculate("age", YEAR) + months = self.calculate("months_receiving_social_security_disability", YEAR) + return (age >= 65) | (months >= 24) + + # household + def _v_county(self): + if "county" in self.household: + return EnumResult([self.household["county"][YEAR]]) + return EnumResult([FIRST_COUNTY.get(self.state, "UNKNOWN")]) + + # tax unit + def _v_adjusted_gross_income(self): + return np.array( + [self._sum("employment_income") + self._sum("social_security_disability")] + ) + + def _v_state_withheld_income_tax(self): + agi = self._sum("adjusted_gross_income") + rate = 0.05 if self.state in INCOME_TAX_STATES else 0.0 + return self._unit( + self.tax_unit, "state_withheld_income_tax", lambda: rate * agi + ) + + def _v_state_sales_tax(self): + return np.array([800.0]) + + def _v_local_sales_tax(self): + share = 0.0 if self.state in NO_LOCAL_SALES_TAX_STATES else 0.2 + return self._unit( + self.tax_unit, + "local_sales_tax", + lambda: share * self._sum("state_sales_tax"), + ) + + def _v_local_income_tax(self): + # NYC residents owe a toy 3% city tax. + rate = 0.03 if self._sum("in_nyc") else 0.0 + return np.array([rate * self._sum("adjusted_gross_income")]) + + def _v_salt_deduction(self): + income = self._sum("state_withheld_income_tax") + self._sum("local_income_tax") + sales = self._sum("state_sales_tax") + self._sum("local_sales_tax") + real_estate = float(self.tax_unit.get("real_estate_taxes", {YEAR: 0})[YEAR]) + return np.array([min(self.salt_cap, max(income, sales) + real_estate)]) + + def _v_first_home_mortgage_balance(self): + return self._unit(self.tax_unit, "first_home_mortgage_balance", lambda: 0) + + def _v_second_home_mortgage_balance(self): + return self._unit(self.tax_unit, "second_home_mortgage_balance", lambda: 0) + + def _v_first_home_mortgage_origination_year(self): + return self._unit( + self.tax_unit, "first_home_mortgage_origination_year", lambda: 0 + ) + + def _v_second_home_mortgage_origination_year(self): + return self._unit( + self.tax_unit, "second_home_mortgage_origination_year", lambda: 0 + ) + + def _v_deductible_mortgage_interest_tax_unit(self): + interest = float( + self.tax_unit.get("first_home_mortgage_interest", {YEAR: 0})[YEAR] + ) + balance = self._sum("first_home_mortgage_balance") + year = self._sum("first_home_mortgage_origination_year") + cap = 1_000_000.0 if 0 < year <= 2017 else 750_000.0 + share = min(1.0, cap / balance) if balance > 0 else 1.0 + return np.array([interest * share]) + + def _v_non_deductible_mortgage_interest_tax_unit(self): + interest = float( + self.tax_unit.get("first_home_mortgage_interest", {YEAR: 0})[YEAR] + ) + return np.array([interest - self._sum("deductible_mortgage_interest_tax_unit")]) + + def _v_income_tax_before_refundable_credits(self): + # The federal output's engine variable. + agi = self._sum("adjusted_gross_income") + medical = max(0.0, self._sum("medicare_part_b_premium") - 0.075 * agi) + itemized = ( + self._sum("salt_deduction") + + medical + + self._sum("deductible_mortgage_interest_tax_unit") + ) + return np.array([0.2 * max(0.0, agi - max(15_000.0, itemized))]) + + def _v_in_nyc(self): + county = self.calculate("county", YEAR).decode_to_str()[0] + return self._unit(self.household, "in_nyc", lambda: county in NYC_COUNTIES) + + def _v_in_ny_mctd_zone_2(self): + county = self.calculate("county", YEAR).decode_to_str()[0] + return self._unit( + self.household, "in_ny_mctd_zone_2", lambda: county in MCTD_ZONE_2 + ) + + def _v_takes_up_medicare_if_eligible(self): + return self._person("takes_up_medicare_if_eligible", True).astype(bool) + + def _v_home_mortgage_interest_tax_unit(self): + return np.array( + [float(self.tax_unit.get("first_home_mortgage_interest", {YEAR: 0})[YEAR])] + ) + + def _v_md_local_income_tax_before_refundable_credits(self): + if self.state != "MD": + return np.array([0.0]) + county = self.calculate("county", YEAR).decode_to_str()[0] + return np.array([MD_RATES[county] * self._sum("adjusted_gross_income")]) + + def _v_nyc_income_tax_before_refundable_credits(self): + return np.array([0.0]) + + def _v_state_income_tax_before_refundable_credits(self): + if self.state not in INCOME_TAX_STATES: + return np.array([0.0]) + amount = 0.04 * self._sum("adjusted_gross_income") + amount += self.state_salt_share * self._sum("salt_deduction") + if "md_local_income_tax_before_refundable_credits" in self.system.aggregate: + amount += self._sum("md_local_income_tax_before_refundable_credits") + return np.array([max(0.0, amount)]) + + def _v_state_refundable_credits(self): + return self._unit(self.tax_unit, "state_refundable_credits", lambda: 0) + + # benefits + def _v_snap(self): + hours = self.calculate("weekly_hours_worked_before_lsr", YEAR) + ages = self.calculate("age", YEAR) + adults = ages >= 18 + if (ages < 18).any() or (hours[adults] >= 20).all(): + county = self.calculate("county", YEAR).decode_to_str()[0] + return np.array([1_200.0 + SNAP_REGION_BONUS.get(county, 0.0)]) + return np.array([0.0]) + + def _v_ssi(self): + return self.calculate("meets_ssi_disability_criteria", YEAR) * 5_000.0 + + +class DivergingSimulation(FakeSimulation): + """A state tax that rises twice as fast as an uncapped SALT deduction.""" + + state_salt_share = 2.0 + salt_cap = float("inf") + + +@pytest.fixture(autouse=True) +def _reset_counter(): + FakeSimulation.instances = 0 + yield + + +def make_engine(simulation_class=FakeSimulation, **kwargs): + return sweep.Engine( + system=FakeSystem(), + simulation_class=simulation_class, + **kwargs, + ) + + +def make_scenario( + sid="scenario_test", + state="MD", + income=120_000.0, + *, + head_inputs=None, + spouse=None, + children=(), + tax_unit_inputs=None, +): + adults = [ + Person("head", 45, income, dict(head_inputs or {"is_tax_unit_head": True})) + ] + if spouse is not None: + adults.append(spouse) + return Scenario( + id=sid, + state=state, + filing_status="single" if spouse is None else "joint", + adults=adults, + children=list(children), + tax_unit_inputs=dict(tax_unit_inputs or {}), + year=2026, + ) + + +def make_job(scenario, estimate_ids=None, programs=PROGRAMS, combined=True): + from policybench.spec import expand_programs_for_scenario + + ids = tuple(estimate_ids or sweep.ESTIMATES_BY_ID) + return sweep.HouseholdJob( + scenario_json=json.dumps(scenario_to_dict(scenario)), + variables=tuple(expand_programs_for_scenario(programs, scenario)), + estimate_ids=ids, + combined=combined, + ) + + +def readings_of(result, estimate, reading=None): + return [ + r + for r in result["readings"] + if r["estimate"] == estimate and (reading is None or r["reading"] == reading) + ] + + +# --------------------------------------------------------------------------- +# Override +# --------------------------------------------------------------------------- + + +def base_situation(): + return make_scenario().to_pe_household() + + +def test_override_sets_values_for_the_period_without_mutating_the_input(): + situation = base_situation() + before = copy.deepcopy(situation) + override = sweep.Override( + people={"head": {"medicare_part_b_premium": 0.0}}, + tax_unit={"state_withheld_income_tax": 12.5}, + household={"county": "BALTIMORE_CITY_MD"}, + ) + out = override.apply(situation, YEAR) + assert situation == before + assert out["people"]["head"]["medicare_part_b_premium"] == {YEAR: 0.0} + assert out["tax_units"]["tax_unit"]["state_withheld_income_tax"] == {YEAR: 12.5} + assert out["households"]["household"]["county"] == {YEAR: "BALTIMORE_CITY_MD"} + + +def test_override_rejects_an_unknown_person(): + with pytest.raises(KeyError): + sweep.Override(people={"ghost": {"x": 1}}).apply(base_situation(), YEAR) + + +def test_override_merge_rejects_conflicts_and_keeps_both_sides(): + a = sweep.Override(tax_unit={"local_sales_tax": 0.0}) + b = sweep.Override(tax_unit={"state_withheld_income_tax": 0.0}) + merged = a.merged(b) + assert merged.tax_unit == {"local_sales_tax": 0.0, "state_withheld_income_tax": 0.0} + with pytest.raises(ValueError): + a.merged(sweep.Override(tax_unit={"local_sales_tax": 1.0})) + with pytest.raises(ValueError): + sweep.Override(people={"head": {"x": 1}}).merged( + sweep.Override(people={"head": {"x": 2}}) + ) + + +def test_override_key_is_canonical_across_number_types(): + assert ( + sweep.Override(tax_unit={"x": 0}).key() + == sweep.Override(tax_unit={"x": 0.0}).key() + ) + assert sweep.Override().is_empty() + assert sweep.Override(tax_unit={"x": 1.0}).describe() == "tax_unit.x=1.00" + + +values = st.one_of( + st.floats(min_value=-1e6, max_value=1e6, allow_nan=False), st.booleans() +) +variable_names = st.sampled_from(["a", "b", "c", "d"]) +group_values = st.dictionaries(variable_names, values, max_size=3) + + +@settings(max_examples=60, deadline=None) +@given(tax_unit=group_values, household=group_values, head=group_values) +def test_override_apply_is_pure_and_idempotent(tax_unit, household, head): + situation = base_situation() + before = copy.deepcopy(situation) + override = sweep.Override( + people={"head": head} if head else {}, tax_unit=tax_unit, household=household + ) + once = override.apply(situation, YEAR) + assert situation == before + assert override.apply(once, YEAR) == once + + +@settings(max_examples=60, deadline=None) +@given(a=group_values, b=group_values) +def test_override_merge_is_commutative_when_it_succeeds(a, b): + left = sweep.Override(tax_unit=a) + right = sweep.Override(tax_unit=b) + try: + ab = left.merged(right) + except ValueError: + with pytest.raises(ValueError): + right.merged(left) + return + assert ab.key() == right.merged(left).key() + + +# --------------------------------------------------------------------------- +# Moves +# --------------------------------------------------------------------------- + + +def test_amount_moves_only_past_the_dollar_tolerance(): + assert not sweep.output_moved(100.0, 101.0, binary=False) + assert sweep.output_moved(100.0, 101.01, binary=False) + assert sweep.output_moved(100.0, 98.99, binary=False) + + +def test_flag_moves_only_when_it_flips(): + assert sweep.output_moved(0.0, 1.0, binary=True) + assert not sweep.output_moved(1.0, 1.0, binary=True) + assert not sweep.output_moved(1.0, 0.9, binary=True) + + +amounts = st.floats(min_value=-1e7, max_value=1e7, allow_nan=False) + + +@settings(max_examples=200, deadline=None) +@given(a=amounts, b=amounts, binary=st.booleans()) +def test_moves_are_symmetric_and_never_reflexive(a, b, binary): + assert sweep.output_moved(a, b, binary=binary) == sweep.output_moved( + b, a, binary=binary + ) + assert not sweep.output_moved(a, a, binary=binary) + if not binary: + assert sweep.output_moved(a, b, binary=False) == (abs(a - b) > 1.0) + + +def test_binary_outputs_follow_the_spec_metric_type(): + assert sweep.is_binary_output("head_medicare_eligible") + assert sweep.is_binary_output("free_school_meals_eligible") + assert not sweep.is_binary_output("snap") + + +# --------------------------------------------------------------------------- +# Registry +# --------------------------------------------------------------------------- + +REQUIRED = { + "state_withheld_income_tax", + "local_sales_tax", + "medicare_part_b_premium", + "county", + "weekly_hours_worked_before_lsr", + "mortgage_origination_year", +} + + +def test_registry_holds_the_audited_estimates_with_unique_readings(): + ids = [e.id for e in sweep.ESTIMATES] + assert len(ids) == len(set(ids)) + assert REQUIRED <= set(ids) + assert sweep.COMBINED_ESTIMATE not in ids + for estimate in sweep.ESTIMATES: + assert estimate.readings, estimate.id + reading_ids = [r.id for r in estimate.readings] + assert len(reading_ids) == len(set(reading_ids)) + for reading in estimate.readings: + assert reading.kind in {sweep.LITERAL, sweep.ALTERNATIVE} + assert reading.description + for text in ( + estimate.engine_behavior, + estimate.why_unlisted, + estimate.found_in, + ): + assert text.strip() + assert estimate.entity in {"person", "tax_unit", "spm_unit", "household"} + + +def test_state_withholding_has_the_literal_and_both_paid_readings(): + estimate = sweep.ESTIMATES_BY_ID["state_withheld_income_tax"] + kinds = {r.id: r.kind for r in estimate.readings} + assert kinds == { + "zero": sweep.LITERAL, + "liability": sweep.ALTERNATIVE, + "net": sweep.ALTERNATIVE, + } + + +def test_no_frozen_household_states_a_registered_input(): + run = Path(__file__).resolve().parents[1] / sweep.DEFAULT_RUN_DIR + scenarios = pd.read_csv(run / "scenarios.csv") + registered = { + name for estimate in sweep.ESTIMATES for name in estimate.engine_inputs + } + for raw in scenarios["scenario_json"]: + data = json.loads(raw) + stated = set(data["tax_unit_inputs"]) | set(data["household_inputs"]) + stated |= set(data["spm_unit_inputs"]) + for person in data["adults"] + data["children"]: + stated |= set(person["inputs"]) + assert not stated & registered, (data["id"], stated & registered) + + +def test_acknowledged_moves_are_scored_outputs_awaiting_a_ruling(): + """Each acknowledgement names a scored output of the frozen run and the decision + it waits on; once a release excludes the output, the entry must go.""" + root = Path(__file__).resolve().parents[1] + path = root / "reference_audit/unlisted_input_acknowledged.json" + acknowledgements = sweep.load_acknowledgements(path) + assert acknowledgements + run = root / sweep.DEFAULT_RUN_DIR + reference = pd.read_csv(run / "reference_outputs.csv") + outputs = set(zip(reference["scenario_id"], reference["variable"])) + excluded = { + (e["scenario_id"], e["variable"]) + for e in json.loads((run / "reference_exclusions.json").read_text())[ + "exclusions" + ] + } + for (sid, variable, estimate, _), entry in acknowledgements.items(): + assert (sid, variable) in outputs + assert (sid, variable) not in excluded, (sid, variable) + assert estimate == "*" or estimate in sweep.ESTIMATES_BY_ID + assert re.fullmatch(r"pending d\d+", entry["status"]), entry["status"] + + +def test_registered_person_inputs_never_reach_a_prompt(): + for name in ( + "weekly_hours_worked_before_lsr", + "takes_up_medicare_if_eligible", + "county_fips", + ): + assert is_excluded_prompt_input_name(name) + + +def test_select_estimates_rejects_unknown_ids(): + assert sweep.select_estimates(None) == sweep.ESTIMATES + assert [e.id for e in sweep.select_estimates(["county"])] == ["county"] + with pytest.raises(SystemExit): + sweep.select_estimates(["no_such_estimate"]) + + +# --------------------------------------------------------------------------- +# One household on the fake engine +# --------------------------------------------------------------------------- + + +def test_household_baseline_and_withholding_readings_move_federal_tax(): + scenario = make_scenario( + tax_unit_inputs={ + "real_estate_taxes": 4_000.0, + "first_home_mortgage_interest": 20_000.0, + } + ) + result = sweep.sweep_household( + make_job(scenario, ["state_withheld_income_tax"]), make_engine() + ) + base = result["baseline"]["federal_income_tax_before_refundable_credits"] + zero = readings_of(result, "state_withheld_income_tax", "zero")[0] + assert not zero["noop"] + assert zero["override"] == {"tax_unit": {"state_withheld_income_tax": 0.0}} + # Itemized deductions: SALT of min(10,000, 6,000 + 4,000) plus 20,000 of + # interest. Without withholding, SALT takes 960 of sales tax plus 4,000. + assert base == pytest.approx(0.2 * (120_000 - 30_000)) + assert zero["outputs"]["federal_income_tax_before_refundable_credits"] == ( + pytest.approx(base + 0.2 * (10_000 - 4_960)) + ) + + +def test_liability_reading_reaches_a_fixed_point_with_county_tax_added(): + scenario = make_scenario() + engine = make_engine() + result = sweep.sweep_household( + make_job(scenario, ["state_withheld_income_tax"]), engine + ) + reading = readings_of(result, "state_withheld_income_tax", "liability")[0] + assert reading["converged"] + last = reading["trace"][-1] + assert abs(last["implied"] - last["withheld"]) < sweep.FIXED_POINT_TOLERANCE + paid = reading["detail"]["paid"] + # Re-simulating at the reported amount reproduces it: state tax at that SALT + # deduction plus Allegany's county tax. + situation = sweep.Override(tax_unit={"state_withheld_income_tax": paid}).apply( + scenario.to_pe_household(), YEAR + ) + sim = FakeSimulation(engine.system, situation) + state = sim._sum("state_income_tax_before_refundable_credits") + county = 0.0303 * 120_000 + assert state + county == pytest.approx(paid, abs=0.01) + assert reading["detail"]["engine_value"] == pytest.approx(0.05 * 120_000) + + +def test_net_reading_subtracts_state_refundable_credits(): + scenario = make_scenario(tax_unit_inputs={"state_refundable_credits": 300.0}) + result = sweep.sweep_household( + make_job(scenario, ["state_withheld_income_tax"]), make_engine() + ) + liability = readings_of(result, "state_withheld_income_tax", "liability")[0] + net = readings_of(result, "state_withheld_income_tax", "net")[0] + assert net["converged"] + paid = net["detail"]["paid"] + assert paid < liability["detail"]["paid"] + situation = sweep.Override(tax_unit={"state_withheld_income_tax": paid}).apply( + scenario.to_pe_household(), YEAR + ) + sim = FakeSimulation(FakeSystem(), situation) + state = sim._sum("state_income_tax_before_refundable_credits") + assert state + 0.0303 * 120_000 - 300.0 == pytest.approx(paid, abs=0.01) + + +def test_a_diverging_fixed_point_is_reported_not_hidden(): + scenario = make_scenario(income=60_000.0) + result = sweep.sweep_household( + make_job(scenario, ["state_withheld_income_tax"]), + make_engine(DivergingSimulation), + ) + reading = readings_of(result, "state_withheld_income_tax", "liability")[0] + assert not reading["converged"] + assert reading["iterations"] == sweep.MAX_FIXED_POINT_ITERATIONS + # The reported amount is the one last simulated, not the next iterate. + assert reading["detail"]["paid"] == reading["trace"][-1]["withheld"] + assert ( + reading["override"]["tax_unit"]["state_withheld_income_tax"] + == (reading["trace"][-1]["withheld"]) + ) + + +def test_readings_that_change_nothing_are_not_simulated(): + # Texas has no income tax: the withholding estimate is already 0, and so is + # the liability, so all three withholding readings are no-ops. + scenario = make_scenario(state="TX") + result = sweep.sweep_household( + make_job(scenario, ["state_withheld_income_tax"], combined=False), + make_engine(), + ) + withholding = readings_of(result, "state_withheld_income_tax") + assert {r["reading"] for r in withholding} == {"zero", "liability", "net"} + assert all(r["noop"] and r["outputs"] is None for r in withholding) + assert result["simulations"] == 1 + + +def test_local_sales_tax_reading_is_a_noop_where_the_engine_already_has_none(): + ct = sweep.sweep_household( + make_job(make_scenario(state="CT"), ["local_sales_tax"]), make_engine() + ) + md = sweep.sweep_household( + make_job(make_scenario(state="MD"), ["local_sales_tax"]), make_engine() + ) + assert readings_of(ct, "local_sales_tax")[0]["noop"] + assert not readings_of(md, "local_sales_tax")[0]["noop"] + + +def test_part_b_reading_zeroes_only_people_with_a_modeled_premium(): + spouse = Person("spouse", 70, 0.0, {"is_tax_unit_spouse": True}) + scenario = make_scenario(spouse=spouse, income=30_000.0) + result = sweep.sweep_household( + make_job(scenario, ["medicare_part_b_premium"]), make_engine() + ) + reading = readings_of(result, "medicare_part_b_premium")[0] + assert reading["override"] == { + "people": {"spouse": {"medicare_part_b_premium": 0.0}} + } + young = sweep.sweep_household( + make_job(make_scenario(), ["medicare_part_b_premium"]), make_engine() + ) + assert readings_of(young, "medicare_part_b_premium") == [] + + +def test_forty_hours_applies_only_to_people_without_stated_hours(): + spouse = Person( + "spouse", + 40, + 10_000.0, + {"is_tax_unit_spouse": True, "hours_worked_last_week": 35}, + ) + scenario = make_scenario(spouse=spouse, income=0.0) + result = sweep.sweep_household( + make_job(scenario, ["weekly_hours_worked_before_lsr"]), make_engine() + ) + reading = readings_of(result, "weekly_hours_worked_before_lsr")[0] + assert reading["override"] == { + "people": {"head": {"weekly_hours_worked_before_lsr": 40.0}} + } + assert result["baseline"]["snap"] == 0.0 + assert reading["outputs"]["snap"] == 1_200.0 + + +def test_pre_tcja_reading_runs_only_where_the_cap_binds(): + capped = make_scenario( + tax_unit_inputs={ + "first_home_mortgage_balance": 900_000.0, + "first_home_mortgage_interest": 45_000.0, + } + ) + under = make_scenario( + tax_unit_inputs={ + "first_home_mortgage_balance": 300_000.0, + "first_home_mortgage_interest": 15_000.0, + } + ) + result = sweep.sweep_household( + make_job(capped, ["mortgage_origination_year"]), make_engine() + ) + reading = readings_of(result, "mortgage_origination_year")[0] + assert reading["override"] == { + "tax_unit": {"first_home_mortgage_origination_year": 2017.0} + } + assert reading["detail"]["capped_interest"] == pytest.approx( + 45_000 * (1 - 750_000 / 900_000) + ) + negative = make_scenario( + tax_unit_inputs={ + "first_home_mortgage_balance": 900_000.0, + "first_home_mortgage_interest": -9_000.0, + } + ) + # Negative interest: nothing is non-deductible, yet the share is below 1, + # so the pre-TCJA cap still changes the deduction and must be simulated. + moved = sweep.sweep_household( + make_job(negative, ["mortgage_origination_year"]), make_engine() + ) + assert len(readings_of(moved, "mortgage_origination_year")) == 1 + base = result["baseline"]["federal_income_tax_before_refundable_credits"] + assert reading["outputs"]["federal_income_tax_before_refundable_credits"] < base + none = sweep.sweep_household( + make_job(under, ["mortgage_origination_year"]), make_engine() + ) + assert readings_of(none, "mortgage_origination_year") == [] + + +def test_county_reading_tries_every_other_county_of_the_state(): + engine = make_engine() + assert engine.counties("MD") == [ + "ALLEGANY_COUNTY_MD", + "BALTIMORE_CITY_MD", + "WORCESTER_COUNTY_MD", + ] + scenario = make_scenario(tax_unit_inputs={"real_estate_taxes": 0.0}) + result = sweep.sweep_household(make_job(scenario, ["county"]), engine) + readings = readings_of(result, "county") + assert [r["variant"] for r in readings] == engine.counties("MD") + by_county = {r["variant"]: r for r in readings} + assert by_county["ALLEGANY_COUNTY_MD"]["noop"] + assert not by_county["WORCESTER_COUNTY_MD"]["noop"] + + +def test_county_states_limits_the_county_reading(): + engine = make_engine(county_states=frozenset({"CA"})) + result = sweep.sweep_household(make_job(make_scenario(), ["county"]), engine) + assert readings_of(result, "county") == [] + + +def test_only_moves_the_locality_alone_explains_are_labeled_prompt_rules_out(): + engine = make_engine() + assert {"in_nyc", "in_ny_mctd_zone_2"} <= set(engine.locality_flags) + child = Person("child1", 8, 0.0, {}) + scenario = make_scenario(state="NY", income=80_000.0, children=[child]) + job = make_job(scenario, ["county"], programs=PROGRAMS + ["local_income_tax"]) + result = sweep.sweep_household(job, engine) + assert result["baseline_localities"] == [] + by_county = {r["variant"]: r for r in readings_of(result, "county")} + kings, nassau = by_county["KINGS_COUNTY_NY"], by_county["NASSAU_COUNTY_NY"] + assert kings["localities"] == ["in_nyc"] + assert nassau["localities"] == ["in_ny_mctd_zone_2"] + # NYC tax moves only through in_nyc; SNAP's region reads the county itself. + assert kings["locality_outputs"] == ["local_income_tax"] + assert nassau["locality_outputs"] == [] + assert kings["outputs"]["local_income_tax"] == 2_400.0 + assert kings["outputs"]["snap"] == 1_500.0 + assert nassau["outputs"]["snap"] == 1_400.0 + reference = pd.Series( + {(scenario.id, v): value for v, value in result["baseline"].items()} + ) + report = sweep.evaluate([result], reference, {}) + moves = report.moves.set_index(["variable", "variant"]) + assert moves.loc[("local_income_tax", "KINGS_COUNTY_NY"), "status"] == ( + "prompt_rules_out" + ) + assert moves.loc[("local_income_tax", "KINGS_COUNTY_NY"), "localities"] == "in_nyc" + assert moves.loc[("snap", "KINGS_COUNTY_NY"), "status"] == "scored" + assert moves.loc[("snap", "NASSAU_COUNTY_NY"), "status"] == "scored" + assert report.summary["counts"]["scored"] == 1 + assert report.summary["counts"]["prompt_rules_out"] == 1 + summary = report.summary + assert sweep.exit_code(summary, strict=True, allow_baseline_mismatch=False) == 2 + + +def test_medicare_not_enrolled_reading_turns_take_up_off_for_eligible_people(): + spouse = Person("spouse", 70, 0.0, {"is_tax_unit_spouse": True}) + scenario = make_scenario(spouse=spouse, income=30_000.0) + result = sweep.sweep_household( + make_job(scenario, ["medicare_part_b_premium"]), make_engine() + ) + zero, not_enrolled = readings_of(result, "medicare_part_b_premium") + assert not_enrolled["reading"] == "not_enrolled" + assert not_enrolled["kind"] == sweep.ALTERNATIVE + assert not_enrolled["override"] == { + "people": {"spouse": {"takes_up_medicare_if_eligible": False}} + } + assert not_enrolled["outputs"] == zero["outputs"] + + +def test_locality_flags_are_discovered_from_the_installed_engine(): + flags = sweep.discover_locality_flags() + assert {"in_nyc", "in_san_francisco", "in_denver"} <= set(flags) + assert all(flag.startswith("in_") for flag in flags) + # Indiana's state variables share the in_ prefix; none is a locality flag. + assert "in_income_tax" not in flags + + +def test_ssdi_and_ssi_alternatives_target_the_right_people(): + head = { + "is_tax_unit_head": True, + "social_security_disability": 12_000.0, + "is_disabled": True, + } + scenario = make_scenario(head_inputs=head, income=0.0) + result = sweep.sweep_household( + make_job( + scenario, + [ + "months_receiving_social_security_disability", + "meets_ssi_disability_criteria", + ], + ), + make_engine(), + ) + months = readings_of(result, "months_receiving_social_security_disability")[0] + assert months["outputs"]["head_medicare_eligible"] == 1.0 + assert result["baseline"]["head_medicare_eligible"] == 0.0 + ssi = readings_of(result, "meets_ssi_disability_criteria")[0] + assert ssi["outputs"]["ssi"] == 5_000.0 + + +def test_combined_reading_applies_every_literal_reading_at_once(): + spouse = Person("spouse", 70, 0.0, {"is_tax_unit_spouse": True}) + scenario = make_scenario(spouse=spouse, income=90_000.0) + result = sweep.sweep_household(make_job(scenario), make_engine()) + combined = readings_of(result, sweep.COMBINED_ESTIMATE)[0] + assert combined["detail"]["parts"] == [ + "state_withheld_income_tax/zero", + "local_sales_tax/zero", + "medicare_part_b_premium/zero", + ] + assert combined["override"]["tax_unit"] == { + "local_sales_tax": 0.0, + "state_withheld_income_tax": 0.0, + } + assert combined["override"]["people"] == { + "spouse": {"medicare_part_b_premium": 0.0} + } + + +def test_combined_reading_is_skipped_when_one_literal_reading_or_none_applies(): + # Texas: no withholding estimate and no Medicare-age person, so only the local + # sales tax reading changes anything and the combined reading would repeat it. + result = sweep.sweep_household(make_job(make_scenario(state="TX")), make_engine()) + assert readings_of(result, sweep.COMBINED_ESTIMATE) == [] + + +def test_identical_overrides_are_simulated_once(): + # Without refundable credits the net reading's fixed point is the liability's. + scenario = make_scenario() + result = sweep.sweep_household( + make_job(scenario, ["state_withheld_income_tax"], combined=False), + make_engine(), + ) + liability = readings_of(result, "state_withheld_income_tax", "liability")[0] + net = readings_of(result, "state_withheld_income_tax", "net")[0] + assert liability["override"] == net["override"] + # baseline + zero + the liability iterations; net reuses every one of them. + assert result["simulations"] == 2 + liability["iterations"] + assert FakeSimulation.instances == result["simulations"] + + +def test_reference_situation_renames_and_patches_as_the_builder_did(): + head = {"is_tax_unit_head": True, "partnership_se_income": 5_000.0} + scenario = make_scenario(head_inputs=head) + engine = make_engine() + + def patch(situation, scenario): + situation["tax_units"]["tax_unit"]["patched"] = {YEAR: 1} + return situation + + situation = sweep.build_reference_situation(scenario, engine, patch) + head_inputs = situation["people"]["head"] + assert "partnership_se_income" not in head_inputs + assert head_inputs["partnership_self_employment_net_earnings"] == {YEAR: 5_000.0} + assert situation["tax_units"]["tax_unit"]["patched"] == {YEAR: 1} + + engine.system.variables["partnership_se_income"] = SimpleNamespace() + kept = sweep.build_reference_situation(scenario, engine) + assert "partnership_se_income" in kept["people"]["head"] + + +def test_household_reports_local_taxes_for_the_scope_check(): + engine = make_engine( + removed_local_components=("md_local_income_tax_before_refundable_credits",) + ) + result = sweep.sweep_household(make_job(make_scenario(), ["county"]), engine) + assert result["local_taxes"] == { + "md_local_income_tax_before_refundable_credits": pytest.approx(0.0303 * 120_000) + } + assert result["removed_local_components"] == [ + "md_local_income_tax_before_refundable_credits" + ] + + +# --------------------------------------------------------------------------- +# SALT paid: two system configurations, one amount +# --------------------------------------------------------------------------- + + +class StubSimulation: + def __init__(self, values): + self.values = values + + def calculate(self, variable, period): + return np.array([self.values.get(variable, 0.0)]) + + +SALT_CONTEXTS = ( + sweep.HouseholdContext(make_scenario(), {}, [], make_engine()), + sweep.HouseholdContext( + make_scenario(), + {}, + [], + make_engine( + aggregate_local_components=( + "md_local_income_tax_before_refundable_credits", + "nyc_income_tax_before_refundable_credits", + ), + remaining_local_components=( + "md_local_income_tax_before_refundable_credits", + "nyc_income_tax_before_refundable_credits", + "nyc_refundable_credits", + ), + ), + ), +) + + +@settings(max_examples=200, deadline=None) +@given( + state=st.floats(min_value=0, max_value=1e5), + md=st.floats(min_value=0, max_value=1e4), + nyc=st.floats(min_value=0, max_value=1e4), + refundable=st.floats(min_value=0, max_value=1e4), + nyc_credits=st.floats(min_value=0, max_value=1e3), + net=st.booleans(), +) +def test_salt_paid_is_the_same_with_or_without_the_scope_adapter( + state, md, nyc, refundable, nyc_credits, net +): + """With the adapter the state lists are state-only; without it they carry the + local taxes and NYC's credits. Either way the amount paid is state tax plus + Maryland county tax, less state refundable credits for the net reading.""" + common = { + "md_local_income_tax_before_refundable_credits": md, + "nyc_income_tax_before_refundable_credits": nyc, + "nyc_refundable_credits": nyc_credits, + } + adapted, raw = SALT_CONTEXTS + with_adapter = sweep.salt_income_tax_paid( + adapted, + StubSimulation( + { + **common, + "state_income_tax_before_refundable_credits": state, + "state_refundable_credits": refundable, + } + ), + net=net, + ) + without = sweep.salt_income_tax_paid( + raw, + StubSimulation( + { + **common, + "state_income_tax_before_refundable_credits": state + md + nyc, + "state_refundable_credits": refundable + nyc_credits, + } + ), + net=net, + ) + expected = max(0.0, state + md - (refundable if net else 0.0)) + assert with_adapter == pytest.approx(expected, abs=1e-6) + assert without == pytest.approx(expected, abs=1e-6) + + +# --------------------------------------------------------------------------- +# Evaluation +# --------------------------------------------------------------------------- + + +def result_for(sid, baseline, readings, local_taxes=None): + return { + "scenario_id": sid, + "state": "MD", + "baseline": baseline, + "local_taxes": local_taxes or {}, + "readings": readings, + "simulations": 1 + len(readings), + "aggregate_local_components": [], + "removed_local_components": [], + "remaining_local_components": [], + "baseline_localities": [], + } + + +def reading( + estimate, + name, + outputs, + kind=sweep.LITERAL, + variant="", + converged=True, + localities=(), + locality_outputs=(), +): + return { + "estimate": estimate, + "reading": name, + "kind": kind, + "variant": variant, + "override": {}, + "override_text": "", + "noop": outputs is None, + "converged": converged, + "iterations": 1, + "trace": [], + "detail": {}, + "localities": list(localities), + "locality_outputs": list(locality_outputs), + "outputs": outputs, + } + + +FED = "federal_income_tax_before_refundable_credits" +STATE = "state_income_tax_before_refundable_credits" + + +def frame(results): + rows = {} + for r in results: + for variable, value in r["baseline"].items(): + rows[(r["scenario_id"], variable)] = value + return pd.Series(rows) + + +def test_evaluate_marks_scored_excluded_and_acknowledged_moves(): + results = [ + result_for( + "s1", + {FED: 1_000.0, STATE: 500.0, "head_medicare_eligible": 0.0}, + [ + reading( + "state_withheld_income_tax", + "zero", + {FED: 1_400.0, STATE: 500.5, "head_medicare_eligible": 0.0}, + ), + reading( + "months_receiving_social_security_disability", + "twenty_four_months", + {FED: 1_000.0, STATE: 500.0, "head_medicare_eligible": 1.0}, + kind=sweep.ALTERNATIVE, + ), + ], + ), + result_for( + "s2", + {FED: 2_000.0}, + [reading("local_sales_tax", "zero", {FED: 2_002.0})], + ), + result_for( + "s3", + {FED: 3_000.0}, + [reading("county", "each_county", {FED: 3_050.0}, variant="X_MD")], + ), + result_for( + "s4", + {FED: 4_000.0}, + [ + reading( + "county", + "each_county", + {FED: 4_100.0}, + variant="KINGS_COUNTY_NY", + localities=["in_nyc"], + locality_outputs=[FED], + ) + ], + ), + ] + exclusions = { + ("s1", "head_medicare_eligible"): { + "reason_code": "reference_depends_on_unlisted_input", + "unlisted_input": "months_receiving_social_security_disability", + }, + ("s2", FED): {"reason_code": "reference_engine_defect"}, + } + acknowledgements = {("s3", FED, "county", "*"): {"status": "pending d999"}} + report = sweep.evaluate(results, frame(results), exclusions, acknowledgements) + moves = report.moves.set_index(["scenario_id", "variable"]) + assert moves.loc[("s1", FED), "status"] == "scored" + assert moves.loc[("s1", FED), "delta"] == pytest.approx(400.0) + assert ("s1", STATE) not in moves.index # 50 cents is within tolerance + assert moves.loc[("s1", "head_medicare_eligible"), "status"] == ( + "excluded_same_input" + ) + assert moves.loc[("s2", FED), "status"] == "excluded_other_reason" + assert moves.loc[("s3", FED), "status"] == "acknowledged" + assert moves.loc[("s3", FED), "acknowledged_status"] == "pending d999" + assert moves.loc[("s4", FED), "status"] == "prompt_rules_out" + assert report.summary["counts"] == { + "scored": 1, + "acknowledged": 1, + "prompt_rules_out": 1, + "excluded_same_input": 1, + "excluded_other_reason": 1, + } + assert ( + sweep.exit_code(report.summary, strict=False, allow_baseline_mismatch=False) + == 0 + ) + assert ( + sweep.exit_code(report.summary, strict=True, allow_baseline_mismatch=False) == 2 + ) + + +def test_evaluate_flags_baseline_mismatches_on_scored_outputs_only(): + results = [result_for("s1", {FED: 1_000.0, STATE: 10.0}, [])] + reference = pd.Series({("s1", FED): 1_000.5, ("s1", STATE): 99.0}) + exclusions = {("s1", STATE): {"reason_code": "reference_engine_defect"}} + report = sweep.evaluate(results, reference, exclusions) + assert report.summary["baseline_scored_mismatches"] == [ + { + "scenario_id": "s1", + "variable": FED, + "reference": 1_000.5, + "baseline": 1_000.0, + } + ] + assert report.summary["baseline_excluded_mismatches"] == 1 + summary = report.summary + assert sweep.exit_code(summary, strict=False, allow_baseline_mismatch=False) == 1 + assert sweep.exit_code(summary, strict=False, allow_baseline_mismatch=True) == 0 + + +def test_scope_violations_and_unconverged_fixed_points_fail_the_run(): + results = [ + result_for( + "s1", + {FED: 1.0}, + [ + reading( + "state_withheld_income_tax", + "liability", + {FED: 1.0}, + converged=False, + ) + ], + local_taxes={"nyc_income_tax_before_refundable_credits": 250.0}, + ) + ] + summary = sweep.evaluate(results, frame(results), {}).summary + assert summary["scope_violations"][0]["scenario_id"] == "s1" + assert summary["unconverged"] == [ + { + "scenario_id": "s1", + "estimate": "state_withheld_income_tax", + "reading": "liability", + } + ] + assert sweep.exit_code(summary, strict=False, allow_baseline_mismatch=True) == 1 + + +def test_exclusion_names_inputs_matches_the_record_text(): + record = {"unlisted_input": "weekly_hours_worked_before_lsr and whether ..."} + assert sweep.exclusion_names_inputs(record, ["weekly_hours_worked_before_lsr"]) + assert not sweep.exclusion_names_inputs(record, ["county"]) + assert not sweep.exclusion_names_inputs(None, ["county"]) + # Whole identifiers only: county_fips does not name county. + fips = {"unlisted_input": "county_fips"} + assert not sweep.exclusion_names_inputs(fips, ["county"]) + assert not sweep.exclusion_names_inputs(record, ["weekly_hours_worked"]) + + +def test_acknowledgements_require_the_key_fields(tmp_path): + path = tmp_path / "ack.json" + path.write_text(json.dumps({"acknowledged": [{"scenario_id": "s1"}]})) + with pytest.raises(SystemExit): + sweep.load_acknowledgements(path) + path.write_text( + json.dumps( + { + "acknowledged": [ + { + "scenario_id": "s1", + "variable": FED, + "estimate": "county", + "status": "pending", + } + ] + } + ) + ) + acknowledgements = sweep.load_acknowledgements(path) + assert ("s1", FED, "county", "*") in acknowledgements + assert sweep.load_acknowledgements(None) == {} + assert sweep._acknowledgement_for(acknowledgements, "s1", FED, "county", "x") + assert not sweep._acknowledgement_for(acknowledgements, "s1", FED, "other", "x") + every = {("s1", FED, "*", "*"): {"status": "pending d963"}} + assert sweep._acknowledgement_for(every, "s1", FED, "other", "x") == { + "status": "pending d963" + } + + +def test_scope_lists_must_agree_across_workers(): + a = result_for("s1", {}, []) + b = dict(a, removed_local_components=["nyc_income_tax_before_refundable_credits"]) + assert sweep.scope_lists([a, a]) == {key: [] for key in sweep.SCOPE_KEYS} + assert sweep.scope_lists([b])["removed_local_components"] == [ + "nyc_income_tax_before_refundable_credits" + ] + with pytest.raises(RuntimeError): + sweep.scope_lists([a, b]) + + +@settings(max_examples=100, deadline=None) +@given( + deltas=st.lists( + st.floats(min_value=-50, max_value=50, allow_nan=False), min_size=1, max_size=6 + ), + excluded=st.lists(st.booleans(), min_size=6, max_size=6), +) +def test_every_move_row_is_a_real_move_with_a_consistent_status(deltas, excluded): + variables = [f"v{i}" for i in range(len(deltas))] + baseline = {v: 100.0 for v in variables} + outputs = {v: 100.0 + d for v, d in zip(variables, deltas)} + results = [ + result_for("s1", baseline, [reading("local_sales_tax", "zero", outputs)]) + ] + exclusions = { + ("s1", v): {"reason_code": "reference_engine_defect"} + for v, flag in zip(variables, excluded) + if flag + } + report = sweep.evaluate(results, frame(results), exclusions) + moved = {v for v, d in zip(variables, deltas) if abs(d) > 1.0} + assert set(report.moves["variable"]) == moved + for _, row in report.moves.iterrows(): + assert row["scored"] == (("s1", row["variable"]) not in exclusions) + assert (row["status"] == "scored") == row["scored"] + assert abs(row["delta"]) > 1.0 + + +# --------------------------------------------------------------------------- +# Jobs, workers and the fix module +# --------------------------------------------------------------------------- + + +def test_build_jobs_puts_the_most_counties_first(monkeypatch): + monkeypatch.setattr( + sweep, "_county_count", lambda state: {"MD": 24, "CA": 58}.get(state, 0) + ) + scenarios = scenario_manifest( + [make_scenario("a", "TX"), make_scenario("b", "MD"), make_scenario("c", "CA")] + ) + reference = pd.Series({(sid, FED): 0.0 for sid in ("a", "b", "c")}) + jobs = sweep.build_jobs( + scenarios, reference, [FED, STATE], sweep.ESTIMATES, county_states=None + ) + ids = [json.loads(j.scenario_json)["id"] for j in jobs] + assert ids == ["c", "b", "a"] + assert all(j.variables == (FED,) for j in jobs) + md_only = sweep.build_jobs( + scenarios, reference, [FED], sweep.ESTIMATES, county_states=frozenset({"MD"}) + ) + assert [json.loads(j.scenario_json)["id"] for j in md_only] == ["b", "a", "c"] + + +def test_run_jobs_builds_one_system_per_worker(monkeypatch): + calls = {} + + class Done: + def __init__(self, value): + self.value = value + + def result(self): + return self.value + + class FakePool: + def __init__(self, max_workers, mp_context, initializer, initargs): + calls.update( + workers=max_workers, + method=mp_context.get_start_method(), + initializer=initializer, + initargs=initargs, + submitted=[], + ) + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def submit(self, fn, job): + calls["submitted"].append((fn, job)) + return Done( + {"scenario_id": job, "state": "MD", "simulations": 1, "seconds": 0} + ) + + monkeypatch.setattr(sweep, "ProcessPoolExecutor", FakePool) + monkeypatch.setattr(sweep, "as_completed", lambda futures: reversed(futures)) + seen = [] + config = sweep.EngineConfig(fix_path=None) + out = sweep.run_jobs( + ["j1", "j2"], config, workers=3, progress=lambda d, t, r: seen.append((d, t)) + ) + assert [r["scenario_id"] for r in out] == ["j2", "j1"] + assert seen == [(1, 2), (2, 2)] + assert calls["workers"] == 3 + assert calls["method"] == "spawn" + assert calls["initializer"] is sweep._init_worker + assert calls["initargs"] == (config,) + assert [fn for fn, _ in calls["submitted"]] == [sweep._run_job] * 2 + + +def test_run_jobs_in_process_uses_the_same_worker_setup(monkeypatch): + built = [] + + def fake_build(config): + built.append(config) + return make_engine(), None + + monkeypatch.setattr(sweep, "build_engine", fake_build) + job = make_job(make_scenario(), ["local_sales_tax"]) + out = sweep.run_jobs([job], sweep.EngineConfig(fix_path=None), workers=1) + assert len(built) == 1 + assert out[0]["scenario_id"] == "scenario_test" + + +def test_fix_dir_and_module_load_with_their_support_files(tmp_path): + fixes = tmp_path / "fixes" + fixes.mkdir() + (fixes / "part.py").write_text("VALUE = 7\n") + (fixes / "notes.txt").write_text("not copied\n") + (fixes / "main_fix.py").write_text( + "import importlib.util, json\n" + "from pathlib import Path\n" + "spec = importlib.util.spec_from_file_location(" + "'part', Path(__file__).with_name('part.py'))\n" + "part = importlib.util.module_from_spec(spec)\n" + "spec.loader.exec_module(part)\n" + "TABLE = json.loads(Path(__file__).with_name('table.json').read_text())\n" + "reform = None\n" + "def patch(situation, scenario):\n" + " return situation\n" + ) + support = tmp_path / "elsewhere" / "table.json" + support.parent.mkdir() + support.write_text(json.dumps({"rate": 0.2})) + path = sweep.assemble_fix_dir(fixes / "main_fix.py", [support], tmp_path / "out") + assert sorted(p.name for p in path.parent.iterdir()) == [ + "main_fix.py", + "part.py", + "table.json", + ] + module = sweep.load_fix_module(path) + assert module.part.VALUE == 7 + assert module.TABLE == {"rate": 0.2} + assert module.reform is None and callable(module.patch) + + +def test_composed_reform_runs_the_fix_then_removes_local_taxes(): + applied = [] + + class FixReform: + @staticmethod + def apply(system): + applied.append("fix") + + class FakeReformSystem: + def __init__(self): + self.parameters = fake_parameters( + [ + "md_income_tax_before_refundable_credits", + "ny_income_tax_before_refundable_credits", + "nyc_income_tax_before_refundable_credits", + ] + ) + + def modify_parameters(self, fn): + self.parameters = fn(self.parameters) + + composed, seen = sweep.compose_reform(FixReform, adapter=True) + system = FakeReformSystem() + composed.apply(system) + # A second pass sees the scoped lists; the record keeps what the fix left. + composed.apply(system) + assert applied == ["fix", "fix"] + assert seen["listed"] == ["nyc_income_tax_before_refundable_credits"] + aggregate = system.parameters.gov.states.household + assert aggregate.state_income_tax_before_refundable_credits("2026-01-01") == [ + "md_income_tax_before_refundable_credits", + "ny_income_tax_before_refundable_credits", + ] + + unadapted, seen = sweep.compose_reform(None, adapter=False) + system = FakeReformSystem() + unadapted.apply(system) + assert seen["listed"] == ["nyc_income_tax_before_refundable_credits"] + assert "nyc_income_tax_before_refundable_credits" in ( + system.parameters.gov.states.household.state_income_tax_before_refundable_credits( + "2026-01-01" + ) + ) + + +# --------------------------------------------------------------------------- +# The command +# --------------------------------------------------------------------------- + + +def write_run(tmp_path, scenarios, engine, exclusions=()): + run = tmp_path / "run" + run.mkdir() + scenario_manifest(scenarios).to_csv(run / "scenarios.csv", index=False) + rows = [] + for scenario in scenarios: + job = make_job(scenario, [], combined=False) + result = sweep.sweep_household(job, engine) + for variable, value in result["baseline"].items(): + rows.append( + {"scenario_id": scenario.id, "variable": variable, "value": value} + ) + pd.DataFrame(rows).to_csv(run / "reference_outputs.csv", index=False) + (run / "reference_outputs.csv.meta.json").write_text( + json.dumps({"programs": PROGRAMS}) + ) + (run / "reference_exclusions.json").write_text( + json.dumps({"exclusions": list(exclusions)}) + ) + return run + + +def test_command_writes_every_artifact_and_gates_on_scored_moves( + tmp_path, monkeypatch, capsys +): + from policybench.cli import main + + engine = make_engine() + scenarios = [ + make_scenario( + "scenario_a", "MD", tax_unit_inputs={"real_estate_taxes": 4_000.0} + ), + make_scenario("scenario_b", "TX", income=40_000.0), + ] + run = write_run(tmp_path, scenarios, engine) + monkeypatch.setattr(sweep, "build_engine", lambda config: (engine, None)) + monkeypatch.setattr(sweep, "_county_count", lambda state: 3) + out = tmp_path / "out" + argv = [ + "policybench", + "unlisted-input-sweep", + "--run-dir", + str(run), + "--out-dir", + str(out), + "--workers", + "1", + ] + monkeypatch.setattr("sys.argv", argv) + main() + for name in ( + "moves.csv", + "readings.csv", + "baseline.csv", + "values.csv.gz", + "summary.json", + "households.json", + "report.md", + ): + assert (out / name).exists(), name + summary = json.loads((out / "summary.json").read_text()) + assert summary["baseline_scored_mismatches"] == [] + assert summary["counts"]["scored"] >= 1 + assert summary["run"]["fix_id"] == "baseline" + assert {e["id"] for e in summary["estimates"]} == set(sweep.ESTIMATES_BY_ID) + moves = pd.read_csv(out / "moves.csv") + # Texas has no withholding estimate; both households have no stated hours. + assert set( + moves[moves["estimate"] == "state_withheld_income_tax"]["scenario_id"] + ) == {"scenario_a"} + assert set( + moves[moves["estimate"] == "weekly_hours_worked_before_lsr"]["variable"] + ) == {"snap"} + assert "Scored outputs that move" in capsys.readouterr().out + + monkeypatch.setattr("sys.argv", argv + ["--strict"]) + with pytest.raises(SystemExit) as raised: + main() + assert raised.value.code == 2 + + acknowledged = tmp_path / "ack.json" + acknowledged.write_text( + json.dumps( + { + "acknowledged": [ + { + "scenario_id": row.scenario_id, + "variable": row.variable, + "estimate": row.estimate, + "status": "pending", + } + for row in moves.itertuples() + ] + } + ) + ) + monkeypatch.setattr( + "sys.argv", argv + ["--strict", "--acknowledged", str(acknowledged)] + ) + main() + + +def test_command_fails_when_a_reference_row_is_not_recomputed( + tmp_path, monkeypatch, capsys +): + from policybench.cli import main + + engine = make_engine() + run = write_run(tmp_path, [make_scenario("scenario_a", "TX")], engine) + reference = pd.read_csv(run / "reference_outputs.csv") + extra = {"scenario_id": "scenario_a", "variable": "tanf", "value": 0.0} + pd.concat([reference, pd.DataFrame([extra])]).to_csv( + run / "reference_outputs.csv", index=False + ) + monkeypatch.setattr(sweep, "build_engine", lambda config: (engine, None)) + monkeypatch.setattr(sweep, "_county_count", lambda state: 0) + out = tmp_path / "out" + monkeypatch.setattr( + "sys.argv", + [ + "policybench", + "unlisted-input-sweep", + "--run-dir", + str(run), + "--out-dir", + str(out), + "--workers", + "1", + "--estimate", + "local_sales_tax", + ], + ) + with pytest.raises(SystemExit) as raised: + main() + assert raised.value.code == 1 + summary = json.loads((out / "summary.json").read_text()) + assert summary["reference_rows_not_recomputed"] == [["scenario_a", "tanf"]] + assert "Every reference row recomputed: NO" in capsys.readouterr().out + + +def test_command_refuses_scenarios_outside_the_scoped_year(tmp_path, monkeypatch): + from policybench.cli import main + + engine = make_engine() + scenario = make_scenario("scenario_a", "TX") + run = write_run(tmp_path, [scenario], engine) + manifest = pd.read_csv(run / "scenarios.csv") + data = json.loads(manifest.loc[0, "scenario_json"]) + data["year"] = 2025 + manifest.loc[0, "scenario_json"] = json.dumps(data) + manifest.to_csv(run / "scenarios.csv", index=False) + monkeypatch.setattr( + "sys.argv", + [ + "policybench", + "unlisted-input-sweep", + "--run-dir", + str(run), + "--out-dir", + str(tmp_path / "out"), + ], + ) + with pytest.raises(SystemExit, match="Scenario years"): + main() + + +def test_command_rejects_a_run_without_its_reference(tmp_path, monkeypatch): + from policybench.cli import main + + monkeypatch.setattr( + "sys.argv", + [ + "policybench", + "unlisted-input-sweep", + "--run-dir", + str(tmp_path), + "--out-dir", + str(tmp_path / "out"), + ], + ) + with pytest.raises(SystemExit, match="Missing"): + main() + + +# --------------------------------------------------------------------------- +# Real engine (slow) +# --------------------------------------------------------------------------- + +ROOT = Path(__file__).resolve().parents[1] +FIX = ROOT / "reference_audit/2026-09-28/fixes/latest_final.py" +SALES_TABLE = ROOT / "reference_audit/2026-09-22/fixes/r19_irs_sales_tax_2025.json" + + +@pytest.fixture(scope="module") +def real_engine(tmp_path_factory): + fix = sweep.assemble_fix_dir(FIX, [SALES_TABLE], tmp_path_factory.mktemp("fix")) + return sweep.build_engine( + sweep.EngineConfig(fix_path=str(fix), county_states=("MD",)) + ) + + +@pytest.mark.slow +def test_real_engine_reproduces_the_recorded_unlisted_input_exclusions(real_engine): + """On the published reference system the sweep reproduces every scored + reference of these households and finds the recorded weekly-hours and SSDI + exclusions as excluded for the same input.""" + engine, patch = real_engine + assert engine.removed_local_components == ( + "nyc_income_tax_before_refundable_credits", + "nyc_refundable_credits", + "ca_sf_wftc", + ) + assert engine.aggregate_local_components == () + assert engine.remaining_local_components == () + assert {"in_nyc", "in_san_francisco"} <= set(engine.locality_flags) + run = ROOT / sweep.DEFAULT_RUN_DIR + scenarios = pd.read_csv(run / "scenarios.csv") + scenarios = scenarios[ + scenarios["scenario_id"].isin(["scenario_007", "scenario_056", "scenario_068"]) + ] + reference = pd.read_csv(run / "reference_outputs.csv").set_index( + ["scenario_id", "variable"] + )["value"] + meta = json.loads((run / "reference_outputs.csv.meta.json").read_text()) + exclusions = { + (e["scenario_id"], e["variable"]): e + for e in json.loads((run / "reference_exclusions.json").read_text())[ + "exclusions" + ] + } + jobs = sweep.build_jobs( + scenarios, reference, meta["programs"], sweep.ESTIMATES, frozenset({"MD"}) + ) + results = [sweep.sweep_household(job, engine, patch) for job in jobs] + report = sweep.evaluate(results, reference, exclusions) + assert report.summary["baseline_scored_mismatches"] == [] + assert report.summary["scope_violations"] == [] + moved = { + (o["scenario_id"], o["variable"]): o["status"] + for o in report.summary["moved_outputs"] + } + assert moved[("scenario_007", "head_medicare_eligible")] == "excluded_same_input" + assert moved[("scenario_056", "snap")] == "excluded_same_input" + county = [r for r in results if r["scenario_id"] == "scenario_068"][0] + assert len(readings_of(county, "county")) == 24 + + +@pytest.mark.slow +def test_real_engine_keeps_county_rules_scored_inside_flagged_counties(real_engine): + """scenario_118 (NY, SNAP): in NYC and in Nassau or Suffolk the county turns on a + locality flag, but SNAP's utility allowance region reads the county itself, so + the SNAP move is not attributed to the locality.""" + engine, patch = real_engine + ny_engine = sweep.Engine( + system=engine.system, + simulation_class=engine.simulation_class, + aggregate_local_components=engine.aggregate_local_components, + removed_local_components=engine.removed_local_components, + remaining_local_components=engine.remaining_local_components, + county_states=frozenset({"NY"}), + locality_flags=engine.locality_flags, + ) + run = ROOT / sweep.DEFAULT_RUN_DIR + scenarios = pd.read_csv(run / "scenarios.csv") + scenarios = scenarios[scenarios["scenario_id"] == "scenario_118"] + reference = pd.read_csv(run / "reference_outputs.csv").set_index( + ["scenario_id", "variable"] + )["value"] + meta = json.loads((run / "reference_outputs.csv.meta.json").read_text()) + (job,) = sweep.build_jobs( + scenarios, + reference, + meta["programs"], + [sweep.ESTIMATES_BY_ID["county"]], + frozenset({"NY"}), + combined=False, + ) + result = sweep.sweep_household(job, ny_engine, patch) + by_county = {r["variant"]: r for r in readings_of(result, "county")} + kings, nassau = by_county["KINGS_COUNTY_NY"], by_county["NASSAU_COUNTY_NY"] + assert kings["localities"] == ["in_nyc"] + assert nassau["localities"] == ["in_ny_mctd_zone_2"] + base_snap = result["baseline"]["snap"] + assert kings["outputs"]["snap"] > base_snap + 1 + assert nassau["outputs"]["snap"] > base_snap + 1 + assert "snap" not in kings["locality_outputs"] + assert "snap" not in nassau["locality_outputs"] + assert "local_income_tax" in kings["locality_outputs"]