Skip to content

Solve for state weights#466

Merged
donboyd5 merged 10 commits intomasterfrom
solve-for-state-weights
Mar 19, 2026
Merged

Solve for state weights#466
donboyd5 merged 10 commits intomasterfrom
solve-for-state-weights

Conversation

@donboyd5
Copy link
Collaborator

@donboyd5 donboyd5 commented Mar 19, 2026

Summary

Add Clarabel QP solver pipeline for state weight optimization — the solver counterpart to the target preparation in PR #465, which this depends upon.

  • Clarabel constrained QP solver with elastic slack for infeasibility handling, 0.5% tolerance (matching national reweighting)
  • Parallel batch runner with worker-cached TMD data (8 workers solves 51 states in ~4 minutes)
  • Quality report with target accuracy, weight distortion, exhaustion analysis, and bystander variable check
  • Standalone CLI: python -m tmd.areas.solve_weights --scope states --workers 8
  • Parameter sweep utility for tuning solver settings across tax years
  • Lessons learned document for future maintainers

Running the optimizer

From scratch:

make clean
make data
python -m tmd.areas.prepare_targets --scope states
python -m tmd.areas.solve_weights --scope states --workers 8  # adjust workers as appropriate
python -m tmd.areas.quality_report

To re-solve weights only (if TMD data and targets already exist):

python -m tmd.areas.solve_weights --scope states --workers 8
python -m tmd.areas.quality_report

Key design decisions

  • multiplier_max=25: The "multiplier" is the ratio of a record's chosen (optimal) weight for an area divided by the weight it would have if its area had weights proportionate to their population.

    For example, if a record has a national weight of 200 and we reweight it to reflect state A, which has 10% of the nation's population, its proportionate weight would be 20 -- 10% of 200. If this record type is very common in state A, perhaps its optimal weight might be 40. Its multiplier would be 2.0 -- the optimal weight divided by proportionate weight. We limit this multiplier, setting its maximum value to 25 -- the maximum optimal weight for this record would be 500. We based the 25 maximum on a 12-combination parameter sweep (4 × 3 grid) that examines combinations of the multiplier maximum and the relative importance put on weight changes vs. target violations. With the 25 maximum, the extent to which weights were "overused", where a record's sum of weights across areas is larger than its national weight, was relatively low and target misses were minimal.

  • weight_penalty has no effect: Weight_penalty has no effect: The solver balances two goals — keeping weights close to proportional and hitting targets. The weight_penalty parameter controls the relative importance of these goals. Our sweep showed it only increases violations without changing weight structure or exhaustion. The solver reaches the same solution regardless of penalty weight.

  • Single-pass preferred: Two-pass iterative exhaustion limiting was tested but found too aggressive — proportional cap scaling made targets infeasible (8,979 violations). Tighter multiplier_max in a single pass is simpler and more robust.

  • Filing-status counts excluded from $1M+ bin: Dual variable analysis showed these are the only expensive constraints (dual costs 6–8 orders of magnitude above all others). Removing them eliminated virtually all constraint cost.

Bystander variable check

Untargeted variables checked for cross-state aggregation distortion. Most are well-behaved (<1%), but a few show >2% distortion: student loan interest (-10.5%), AMT (-10.3%), tax-exempt interest (+7.4%), qualified dividends (-4.1%). These are driven by rare high-income PUF records being over/under-weighted. Documented in lessons with guidance on when to worry.

New files

File Purpose
tmd/areas/create_area_weights_clarabel.py Clarabel QP solver
tmd/areas/batch_weights.py Parallel batch runner
tmd/areas/quality_report.py Cross-state diagnostics
tmd/areas/solve_weights.py CLI entry point
tmd/areas/sweep_params.py Parameter grid search
tmd/areas/AREA_WEIGHTING_LESSONS.md Lessons learned
tests/test_solve_weights.py 11 tests

Test plan

  • 11 new tests pass (solver on faux xx area, log parser, scope parsing, area filtering)

@donboyd5 donboyd5 requested a review from martinholmer March 19, 2026 18:45
@donboyd5
Copy link
Collaborator Author

@martinholmer, I am away tomorrow and the next day but can address any comments or questions you have after that. The documentation for this draft PR still needs work.

@martinholmer
Copy link
Collaborator

@donboyd5, Here are the results on my computer after pulling the branch for PR #465 and then the branch for PR #466:

================ 69 passed, 6 skipped in 51.10s ===============================
make test  406.99s user 17.36s system 268% cpu 2:37.97 total
(base) TMD> python -m tmd.areas.prepare_targets --scope states
Preparing state targets (SOI year 2022)...
  Building extended targets for 51 areas...
  Wrote 51 state target files (178 targets each, 3.3s)
(base) TMD> python -m tmd.areas.solve_weights --scope states --workers 4
Pass 1: solving state weights...
Processing 51 areas (up to 178 targets each) with 4 workers...
(Areas shown in completion order, which varies with parallel workers.)

   1  al az ak ar ca co ct dc fl de  [57s elapsed]
  11  ga hi il id ia in ks ky la ma  [98s elapsed]
  21  md mi me mn mo ms mt nc nd nj  [153s elapsed]
  31  ne nh nm ny oh nv ok pa or ri  [195s elapsed]
  41  sc tx sd tn ut va wa vt wi wv  [250s elapsed]
  51  wy  [250s elapsed]

Completed 51/51 areas in 250.4s
17 areas had violated targets (35 targets total). Largest violation: 0.50%.
Run: python -m tmd.areas.quality_report for full details.
Total solve time: 250.4s
(base) TMD> python -m tmd.areas.quality_report
================================================================================
CROSS-STATE QUALITY SUMMARY REPORT
================================================================================

States: 51
Solved: 51
Failed: 0
States with violated targets: 17/51
Total targets: 51 states × 178 = 9124
Total violated targets: 35

TARGET ACCURACY:
  Per-state mean error: avg across states=0.0047, worst state=0.0048
  Per-state max error:  avg across states=0.0050, worst state=0.0050
  Hit rate:  avg=99.6%, min=97.8% (out of 178 targets, tolerance: +/-0.5% + eps)

WEIGHT DISTORTION (multiplier from 1.0):
  RMSE:   avg=0.609, max=2.014
  Min:    avg=0.000, min=0.000
  P05:    avg=0.091, min=0.000
  Median: avg=0.964, range=[0.824, 1.036]
  P95:    avg=1.650, max=3.237
  Max:    avg=21.7, max=25.0

NEAR-ZERO WEIGHT MULTIPLIERS (% of records):
  Exact zero (x=0):  avg=7.8%, max=27.7%
  Below 0.1 (x<0.1): avg=9.0%, max=29.4%

PER-STATE DETAIL:
  Err cols = |relative error| (fraction); weight cols = multiplier on national weight (1.0 = unchanged)
St   Status           Hit   Tot  Viol  MeanErr   MaxErr   wRMSE    wP05    wMed    wP95     wMax  %zero
-------------------------------------------------------------------------------------------------------
AK   Solved           175   178     3   0.0048   0.0050   0.683   0.000   0.933   1.741     25.0  13.4%
AL   Solved           179   179     0   0.0047   0.0050   0.401   0.000   0.911   1.376      9.2   5.0%
AR   Solved           179   179     0   0.0048   0.0050   0.548   0.000   0.897   1.406     25.0   8.0%
AZ   Solved           179   179     0   0.0046   0.0050   0.311   0.416   0.981   1.223     22.8   1.9%
CA   Solved           179   179     0   0.0046   0.0050   0.561   0.000   1.008   1.919     24.8   6.6%
CO   Solved           179   179     0   0.0046   0.0050   0.374   0.523   1.010   1.592     14.7   0.9%
CT   Solved           179   179     0   0.0047   0.0050   0.954   0.000   0.998   2.459     25.0   5.9%
DC   Solved           174   178     4   0.0048   0.0050   1.219   0.000   1.012   3.237     25.0   9.9%
DE   Solved           177   178     1   0.0047   0.0050   0.560   0.000   0.972   1.730     20.2  10.0%
FL   Solved           179   179     0   0.0047   0.0050   0.663   0.000   1.036   1.841     25.0   7.2%
GA   Solved           179   179     0   0.0047   0.0050   0.281   0.487   0.961   1.326      8.3   1.0%
HI   Solved           177   179     2   0.0046   0.0050   0.741   0.000   0.958   1.553     25.0  13.2%
IA   Solved           179   179     0   0.0048   0.0050   0.614   0.000   0.946   1.485     25.0  11.3%
ID   Solved           179   179     0   0.0048   0.0050   0.393   0.156   0.943   1.454     10.6   3.3%
IL   Solved           179   179     0   0.0046   0.0050   0.545   0.178   0.991   1.593     25.0   3.7%
IN   Solved           179   179     0   0.0047   0.0050   0.479   0.000   0.957   1.399     25.0   7.5%
KS   Solved           178   179     1   0.0048   0.0050   0.546   0.000   0.969   1.417     25.0   7.7%
KY   Solved           179   179     0   0.0048   0.0050   0.385   0.000   0.933   1.279     17.4   5.4%
LA   Solved           179   179     0   0.0046   0.0050   0.554   0.000   0.881   1.552     25.0   8.3%
MA   Solved           179   179     0   0.0046   0.0050   0.698   0.242   1.022   2.207     25.0   3.0%
MD   Solved           179   179     0   0.0047   0.0050   0.597   0.000   0.964   1.945     21.6   7.5%
ME   Solved           177   179     2   0.0047   0.0050   0.757   0.000   0.968   1.502     25.0  15.4%
MI   Solved           179   179     0   0.0045   0.0050   0.410   0.066   0.979   1.382     23.4   4.4%
MN   Solved           179   179     0   0.0047   0.0050   0.438   0.258   0.986   1.568     22.4   2.8%
MO   Solved           179   179     0   0.0047   0.0050   0.340   0.223   0.985   1.288     15.3   3.0%
MS   Solved           177   179     2   0.0047   0.0050   0.669   0.000   0.824   1.659     25.0  15.3%
MT   Solved           177   179     2   0.0048   0.0050   0.545   0.000   1.001   1.624     25.0   7.9%
NC   Solved           179   179     0   0.0047   0.0050   0.275   0.410   0.976   1.257     10.8   1.6%
ND   Solved           176   178     2   0.0047   0.0050   1.045   0.000   0.899   2.097     25.0  23.6%
NE   Solved           177   179     2   0.0048   0.0050   0.626   0.000   0.961   1.553     25.0   9.2%
NH   Solved           178   179     1   0.0047   0.0050   0.842   0.000   0.977   2.244     25.0   9.0%
NJ   Solved           179   179     0   0.0047   0.0050   0.722   0.006   0.971   2.147     25.0   4.8%
NM   Solved           177   179     2   0.0048   0.0050   0.565   0.000   0.939   1.331     25.0  11.5%
NV   Solved           179   179     0   0.0046   0.0050   0.592   0.000   1.006   1.652     25.0   5.7%
NY   Solved           179   179     0   0.0047   0.0050   1.005   0.000   0.993   2.193     25.0  11.2%
OH   Solved           179   179     0   0.0047   0.0050   0.515   0.000   0.983   1.345     25.0   7.4%
OK   Solved           179   179     0   0.0047   0.0050   0.400   0.000   0.923   1.320     13.3   5.2%
OR   Solved           179   179     0   0.0046   0.0050   0.359   0.246   0.982   1.400      8.3   2.4%
PA   Solved           179   179     0   0.0046   0.0050   0.331   0.423   0.992   1.391     13.9   1.7%
RI   Solved           178   179     1   0.0048   0.0050   0.488   0.014   1.000   1.478     25.0   4.8%
SC   Solved           179   179     0   0.0047   0.0050   0.328   0.312   0.986   1.307     12.2   2.6%
SD   Solved           178   179     1   0.0048   0.0050   0.857   0.000   0.961   1.751     25.0  15.7%
TN   Solved           179   179     0   0.0047   0.0050   0.463   0.007   1.000   1.440     25.0   4.6%
TX   Solved           179   179     0   0.0047   0.0050   0.740   0.000   0.947   1.735     25.0   8.7%
UT   Solved           179   179     0   0.0047   0.0050   0.520   0.020   0.968   1.574     25.0   4.6%
VA   Solved           179   179     0   0.0046   0.0050   0.391   0.403   0.978   1.589     13.5   2.1%
VT   Solved           175   178     3   0.0047   0.0050   1.012   0.000   0.948   1.728     25.0  21.9%
WA   Solved           179   179     0   0.0046   0.0050   0.604   0.247   0.982   1.823     25.0   3.3%
WI   Solved           179   179     0   0.0047   0.0050   0.511   0.000   0.967   1.514     25.0   6.5%
WV   Solved           177   179     2   0.0048   0.0050   0.570   0.000   0.876   1.362     25.0  15.0%
WY   Solved           175   179     4   0.0048   0.0050   2.014   0.000   0.913   2.156     25.0  27.7%

VIOLATIONS BY VARIABLE:
  c00100: 35 violations across 17 states

STATES WITH MOST VIOLATIONS:
  WY: 4 violated
  DC: 4 violated
  VT: 3 violated
  AK: 3 violated
  ME: 2 violated
  MS: 2 violated
  MT: 2 violated
  HI: 2 violated
  NE: 2 violated
  NM: 2 violated

WORST 5 AMOUNT VIOLATIONS (by % error):
  (none — all amount targets met)

WORST 5 COUNT VIOLATIONS (by % error):
  DC   0.500% target=       3,392  achieved=       3,409  miss=      17  c00100 returns $1000K+
  MT   0.500% target=       2,370  achieved=       2,358  miss=      12  c00100 returns $1000K+
  HI   0.500% target=       2,013  achieved=       2,003  miss=      10  c00100 returns $1000K+
  ME   0.500% target=       2,013  achieved=       2,003  miss=      10  c00100 returns $1000K+
  NM   0.500% target=       2,125  achieved=       2,115  miss=      10  c00100 returns $1000K+

WEIGHT EXHAUSTION (sum of state weights / national weight):
  A ratio of 1.0 means the record's national weight is fully allocated across 51 states.
  min=0.0000, p1=0.2442, p5=0.4420, p10=0.5805, p25=0.8560, median=0.9969, p75=1.0184, p90=1.1995, p95=1.3987, p99=2.0299, max=16.5662
  Mean: 0.9654, Std: 0.3627
  Over-used (>1.10): 36859 (17.1%)  Under-used (<0.90): 61801 (28.7%)
  Exhaustion > 2x: 2280 records
  Exhaustion > 5x: 165 records
  Exhaustion > 10x: 15 records

MOST EXHAUSTED RECORDS (top 5):
  1. RECID 173115: exh=16.6x, s006=9.3, PUF MFJ, AGI=$2,524,724
     wages=$32,953, int=$505,326, div=$684,234, ptshp=$804,066
     top states: TX=21.0, NY=13.8, CA=9.6 (47 nonzero of 51)
  2. RECID 179118: exh=15.5x, s006=452.9, PUF MFJ, AGI=$1,047,640
     wages=$410,649, int=$1,663, div=$29,113, ptshp=$570,596
     top states: TX=1018.0, FL=754.1, IL=426.6 (46 nonzero of 51)
  3. RECID 190633: exh=13.3x, s006=21.4, PUF MFJ, AGI=$6,845,610
     wages=$0, int=$931,630, div=$923,199, ptshp=$1,475,592
     top states: TX=48.1, FL=33.5, NY=31.5 (48 nonzero of 51)
  4. RECID 198795: exh=12.7x, s006=276.3, PUF Single, AGI=$2,027,070
     wages=$0, int=$13,646, div=$708,766, ptshp=$253,206
     top states: TX=517.4, FL=460.1, CA=411.4 (41 nonzero of 51)
  5. RECID 162442: exh=12.7x, s006=22.6, PUF MFJ, AGI=$6,237,402
     wages=$66,418, int=$235,117, div=$842,362, ptshp=$5,168,424
     top states: TX=50.9, NY=31.4, OH=16.4 (47 nonzero of 51)

CROSS-STATE AGGREGATION vs NATIONAL TOTALS for SELECTED VARIABLES (51 states):
  Variable                               National    Sum-of-States    Diff%
  ------------------------------------------------------------------------
  Returns (s006)                      189,765,349      189,505,123   -0.14%
  AGI (c00100)                   $       14874.8B $       14855.5B   -0.13%
  Wages (e00200)                 $        9729.7B $        9711.2B   -0.19%
  Capital gains (capgains_net)   $        1019.5B $         993.9B   -2.51%
  SALT ded (c18300)              $         153.8B $         153.3B   -0.33%
  Income tax (iitax)             $        2147.4B $        2147.8B   +0.02%

BYSTANDER CHECK: UNTARGETED VARIABLES (51 states):
  Variables NOT directly targeted — distortion from weight adjustments.
  Variable                               National    Sum-of-States    Diff%
  ------------------------------------------------------------------------
  Student loan int (e19200)      $         377.5B $         338.1B  -10.46% ***
  AMT (c09600)                   $           1.1B $           1.0B  -10.25% ***
  Tax-exempt int (e00400)        $          56.4B $          60.6B   +7.39% ***
  Qual dividends (e00650)        $         312.6B $         299.9B   -4.06% ***
  Unemployment (e02300)          $          30.6B $          29.4B   -3.98% ***
  Medical expenses (e17500)      $         170.5B $         165.6B   -2.89% ***
  Total credits (c07100)         $         151.3B $         148.1B   -2.17% ***
  Itemized ded (c04470)          $         661.5B $         665.2B   +0.55%
  Sch C income (e00900)          $         411.5B $         409.3B   -0.53%
  Taxable pensions (e01700)      $         915.3B $         910.8B   -0.49%
  Charitable (c19700)            $         201.5B $         200.7B   -0.41%
  Payroll tax                    $        1342.8B $        1339.4B   -0.26%
  Standard deduction             $        3047.1B $        3039.6B   -0.25%
  Children <17 (n24)                   76,365,307       76,186,309   -0.23%
  IRA distrib (e01400)           $         440.2B $         439.2B   -0.22%
  Total persons (XTOT)                333,976,652      333,279,920   -0.21%
  Sch E net (e02000)             $        1147.0B $        1145.2B   -0.16%
  Mortgage int (c19200)          $         234.6B $         234.3B   -0.14%
  Income tax (iitax)             $        2147.4B $        2147.8B   +0.02%

  *** = 7 variables with >2% aggregation distortion

(base) TMD> 

@martinholmer
Copy link
Collaborator

@donboyd5, See my comment in the PR #466 conversation that contains the results of my running the full state pipeline (using the new code in PRs #465 and #466). If you think the results I got on my computer look good, then please merge PR #465 and then PR #466. Thanks for all the work on this!

@martinholmer martinholmer marked this pull request as ready for review March 19, 2026 21:53
@martinholmer martinholmer marked this pull request as draft March 19, 2026 21:59
@martinholmer
Copy link
Collaborator

@donboyd5, Looks like PR #466 is still a work in progress. One example of its incomplete nature is that the
now obsolete create_area_weights.py is still there (using jax, etc.) while the new code seems to be in the create_area_weights_clarabel.py module.

Can you do this:

  • Remove the old create_area_weights.py` module
  • Rename the new create_area_weights_clarabel.py module as create_area_weights.py

As in the national reweight.py, there is no need to mention the solver package in either a module name or a function name. We don't ever create a module called read_csv_pandas.py. Rather we call it read_csv.py and import pandas at the top of the module.

@martinholmer
Copy link
Collaborator

@donboyd5, Don't merge #466 until you do the cleanup of the old code.

donboyd5 and others added 8 commits March 19, 2026 18:18
Port state weight solver pipeline from state-weights-clarabel branch:
- Clarabel constrained QP solver with elastic slack (0.5% tolerance)
- Parallel batch runner with worker-cached TMD data
- Cross-state quality report (log parsing, weight exhaustion, aggregation)
- Standalone CLI: python -m tmd.areas.solve_weights --scope states --workers 8
- 11 tests covering solver, log parser, scope parsing, and area filtering

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Add weight_penalty parameter to QP solver for controlling
  multiplier-vs-constraint tradeoff
- Add --max-exhaustion flag to solve_weights for iterative
  two-pass exhaustion limiting with per-record multiplier caps
- Enhance quality report with exhaustion record profiles
  (top 5 most exhausted records with taxpayer characteristics)
- Add sweep_params.py for grid search over multiplier_max
  and weight_penalty combinations

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Lower AREA_MULTIPLIER_MAX from 100 to 25 based on parameter sweep:
virtually identical target accuracy (35 vs 33 violations) but 34%
lower max exhaustion (16.6x vs 25.2x). Single-pass, no complexity.

Add AREA_WEIGHTING_LESSONS.md documenting parameter tuning findings,
weight exhaustion mechanics, dual variable analysis, SALT targeting,
and guidance for future Congressional district work.

Update README.md with solver usage, quality report, and link to
lessons document.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Quality report now checks ~19 untargeted variables for cross-state
aggregation distortion, sorted by severity, with >2% flagged.

Key bystanders: student loan interest (-10.5%), AMT (-10.3%),
tax-exempt interest (+7.4%), qualified dividends (-4.1%).

Add corresponding section to AREA_WEIGHTING_LESSONS.md explaining
what drives bystander distortion and when to worry about it.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Remove create_area_weights_clarabel.py and replace the old
create_area_weights.py (scipy L-BFGS-B + JAX) with the Clarabel
constrained QP solver. Update all imports and function references.

- Renamed create_area_weights_file_clarabel() to create_area_weights_file()
- Removed valid_area() dependency (areas validated by target file existence)
- Updated imports in solve_weights, batch_weights, quality_report,
  sweep_params, make_all, and test_solve_weights

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@donboyd5
Copy link
Collaborator Author

donboyd5 commented Mar 19, 2026 via email

donboyd5 and others added 2 commits March 19, 2026 18:31
test_area_weights.py tested the old scipy L-BFGS-B solver which is
now replaced by Clarabel. The Clarabel solver is tested by
test_solve_weights.py (test_clarabel_solver_xx and 10 other tests).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Rename test_clarabel_solver_xx to test_solver_xx since the module
name no longer contains "clarabel".

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@donboyd5 donboyd5 force-pushed the solve-for-state-weights branch from fcc479c to d16a9a9 Compare March 19, 2026 22:51
@donboyd5 donboyd5 marked this pull request as ready for review March 19, 2026 22:52
@donboyd5 donboyd5 merged commit a625d39 into master Mar 19, 2026
1 check passed
@donboyd5 donboyd5 deleted the solve-for-state-weights branch March 19, 2026 22:53
@donboyd5
Copy link
Collaborator Author

@martinholmer, thanks for speedy review and helpful comments. The state weights quality reports on your machine was identical to the report on my machine, to every last number. It's nice to see that reproducibility held up with subnational weights also (as we would expect given that it's the same solver).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants