Fix J=1 shape collapse bugs (#1143)

[SeaCelo]

Fixes #1143. Closes #1146.

OG-Core is almost always run with several income groups, and a handful of spots assumed there was always more than one. Each used np.squeeze, which drops a size-1 axis — a no-op for J≥2, but for J=1 it collapses the (S, 1) ability column to a flat (S,). When that flat array is then multiplied against an (S, 1) array, numpy builds an (S, S) outer product instead of multiplying element by element, so the result comes out S× too large or the wrong shape. For J=1 this broke aggregate labor, the income-scaling factor, tax rates, and before-tax income — enough that the steady state couldn't clear (K explodes, r and consumption go negative).

Sites fixed:

• aggregates.get_L — SS labor aggregate
• SS.py average_income_model — calibration factor
• txfunc.get_tax_rates — all parametric tax functions (linear, DEP, DEP_totalinc, GS, HSV), via one shape-aware coefficient helper
• household.get_y — before-tax income
• SS.py — the e argument to the income-tax and pension calls

Why it took a couple of passes: the J=1 test uses a flat productivity profile and the simplest (linear) tax function — settings that happen to hide this bug (an inflated table of ones is still ones), and it never exercises the realistic tax functions. So the test passing wasn't enough; the realistic-tax and before-tax-income sites only surfaced in a full model run. The test also compared against a saved fixture generated under the old buggy code, so it kept failing even after the code was right. That fixture is regenerated.

Validation:

• Full SS + TPI with one income group vs an identical two-group economy (which must give the same answer): match to machine precision across baseline, reform, and use_zeta.
• test_run_TPI_extra[J=1] passes; full suite green.

Regression tests:

• test_get_L_J1_regression — get_L against a plain double-sum reference.
• test_get_tax_rates_J1_shape — every parametric tax function preserves the J=1 shape (fails without the fix).
• test_get_y strengthened to use realistic (S, J) inputs and assert the shape.

Regenerated inner_loop_outputs_J1.pkl and run_TPI_outputs_J1.pkl, which encoded the old buggy output.

cc: @jdebacker @rickecon @utsinboots

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.41%. Comparing base (463156e) to head (b25ca89).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1145      +/-   ##
==========================================
- Coverage   73.51%   73.41%   -0.11%     
==========================================
  Files          21       21              
  Lines        5302     5303       +1     
==========================================
- Hits         3898     3893       -5     
- Misses       1404     1410       +6     

Flag	Coverage Δ
unittests	73.41% <100.00%> (-0.11%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
ogcore/SS.py	79.02% <ø> (-1.99%)	⬇️
ogcore/aggregates.py	100.00% <100.00%> (ø)
ogcore/household.py	90.77% <100.00%> (ø)
ogcore/txfunc.py	46.56% <100.00%> (+0.33%)	⬆️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[utsinboots]

Tested PR #1145, commit 0ca9da89 on Windows

Passed:
uv run pytest tests/test_aggregates.py::test_get_L_J1_regression -vv
uv run pytest tests/test_txfunc.py::test_get_tax_rates_linear_J1_shape -vv
uv run pytest "tests/test_SS.py::test_inner_loop[J=1]" -vv
uv run pytest tests/test_aggregates.py 42 passed

Failed:
uv run pytest "tests/test_TPI.py::test_run_TPI_extra[J=1]" -vv 1 failed in 34.39s
Error: FAILED tests/test_TPI.py::test_run_TPI_extra[J=1] - assert False where False = <function allclose at 0x00000299CE1BBBF0>(array([0.76815272, 0.76666378, 0.77806494, 0.77931536, 0.78144139,\n 0.78393469, 0.78673022, 0.78957472, 0.79246368, 0.79573516, <- - - - - -> \n 0.32000643, 0.32000303, 0.3199997 , 0.31999642, 0.31999319]), rtol=0.0001, atol=0.0001) where <function allclose at 0x00000299CE1BBBF0> = np.allclose

Also, uv run pytest tests/test_txfunc.py 30 passed, 1 failed
Failure was tests/test_txfunc.py::test_tax_func_estimate while reading tests/test_io_data/tax_func_estimate_inputs.pkl with error: ValueError: could not convert string to int

The remaining test_run_TPI_extra[J=1] failure may be related to additional J=1 shape-collapse cases that were intentionally left out of scope for #1145

[jdebacker]

@SeaCelo Thanks for looking into this. I think the code changes you suggest make sense, but I'm still getting the failure identified in Issue #1143:

uv run pytest tests/test_TPI.py::test
_run_TPI_extra

Returns:

============================= test session starts ==============================
platform darwin -- Python 3.11.13, pytest-9.0.2, pluggy-1.6.0
rootdir: /Users/jason.debacker/repos/OG-Core
configfile: pyproject.toml
plugins: anyio-4.12.1, xdist-3.8.0, cov-7.0.0
collected 8 items

tests/test_TPI.py .......F                                               [100%]

=================================== FAILURES ===================================
___________________________ test_run_TPI_extra[J=1] ____________________________
...
=========================== short test summary info ============================
FAILED tests/test_TPI.py::test_run_TPI_extra[J=1] - assert False

[SeaCelo]

@jdebacker @utsinboots Thanks both for checking this.
The first pass left two gaps. The saved comparison file the test checks against was generated under the old buggy code, so the test was comparing correct new numbers against old wrong ones (regenerated now). And the same bug lived in a few spots the test doesn't exercise (the realistic tax functions and before-tax income) which only showed up in a full run.

Fixed now. I ran the full model (SS + TPI) for one income group against an identical two-group economy and they match. I also confirmed the steady state matches across reform and use_zeta. test_run_TPI_extra[J=1] passes and the full suite is green.

Worth confirming on your end with a full J=1 run. the test alone uses simplified settings that hide this class of bug.