Skip to content

Allow model to work with J=1 #1062

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rickecon merged 17 commits into from
Feb 2, 2026
Merged

Allow model to work with J=1 #1062

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
355749e
changes to allow J=1
jdebacker Aug 3, 2025
b86dcaf
Merge remote-tracking branch 'upstream/master' into J1
jdebacker Sep 14, 2025
dee2526
add SS test for J=1
jdebacker Sep 21, 2025
d26e984
add SS test for J=1
jdebacker Sep 28, 2025
d11289c
specify tpi test of j=1
jdebacker Sep 28, 2025
3803c7e
add new test pkl
jdebacker Sep 28, 2025
d1c505b
add test ffile for tpi extra test
jdebacker Sep 29, 2025
7f01b95
fix for rates
jdebacker Sep 29, 2025
9754523
updat efor new saved files
jdebacker Sep 29, 2025
f3dc1d7
more concise reshape
jdebacker Sep 30, 2025
0d20376
update tests to note J=2
jdebacker Sep 30, 2025
005ee38
sync with master
jdebacker Jan 24, 2026
51b18ac
fix spec in test pensions
jdebacker Jan 24, 2026
d4935b3
add filer to spec
jdebacker Jan 24, 2026
67131da
uncomment tests
jdebacker Jan 25, 2026
943465f
update test data
jdebacker Jan 25, 2026
30450ff
bump version and update changelog
jdebacker Jan 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).

## [0.15.4] - 2025-01-27 12:00:00

### Added

- Ability to simulate the model with a single type of household (`J=1`). See PR [#1062](https://github.com/PSLmodels/OG-Core/pull/1062)

## [0.15.3] - 2025-01-24 12:00:00

### Added
Expand Down Expand Up @@ -505,6 +511,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- Any earlier versions of OG-USA can be found in the [`OG-Core`](https://github.com/PSLmodels/OG-Core) repository [release history](https://github.com/PSLmodels/OG-Core/releases) from [v.0.6.4](https://github.com/PSLmodels/OG-Core/releases/tag/v0.6.4) (Jul. 20, 2021) or earlier.


[0.15.4]: https://github.com/PSLmodels/OG-Core/compare/v0.15.3...v0.15.4
[0.15.3]: https://github.com/PSLmodels/OG-Core/compare/v0.15.2...v0.15.3
[0.15.2]: https://github.com/PSLmodels/OG-Core/compare/v0.15.1...v0.15.2
[0.15.1]: https://github.com/PSLmodels/OG-Core/compare/v0.15.0...v0.15.1
Expand Down
22 changes: 11 additions & 11 deletions ogcore/SS.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -418,7 +418,7 @@ def inner_loop(outer_loop_vars, p, client):
None,
False,
"SS",
np.squeeze(p.e[-1, :, :]),
np.squeeze(p.e[-1, :, :]).reshape((p.S, p.J)),
etr_params_3D,
p,
)
Expand All @@ -433,7 +433,7 @@ def inner_loop(outer_loop_vars, p, client):
bq,
rm,
net_tax,
np.squeeze(p.e[-1, :, :]),
np.squeeze(p.e[-1, :, :]).reshape((p.S, p.J)),
p.tau_c[-1, :],
p,
)
Expand Down Expand Up @@ -557,7 +557,7 @@ def inner_loop(outer_loop_vars, p, client):
None,
False,
"SS",
np.squeeze(p.e[-1, :, :]),
np.squeeze(p.e[-1, :, :]).reshape((p.S, p.J)),
etr_params_3D,
p,
)
Expand All @@ -572,7 +572,7 @@ def inner_loop(outer_loop_vars, p, client):
new_bq,
new_rm,
taxss,
np.squeeze(p.e[-1, :, :]),
np.squeeze(p.e[-1, :, :]).reshape((p.S, p.J)),
p.tau_c[-1, :],
p,
)
Expand Down Expand Up @@ -602,7 +602,7 @@ def inner_loop(outer_loop_vars, p, client):
ubi,
theta,
etr_params_3D,
np.squeeze(p.e[-1, :, :]),
np.squeeze(p.e[-1, :, :]).reshape((p.S, p.J)),
p,
None,
"SS",
Expand Down Expand Up @@ -925,7 +925,7 @@ def SS_solver(
nssmat,
factor,
True,
np.squeeze(p.e[-1, :, :]),
np.squeeze(p.e[-1, :, :]).reshape((p.S, p.J)),
etr_params_3D,
mtry_params_3D,
capital_noncompliance_rate_2D,
Expand All @@ -939,7 +939,7 @@ def SS_solver(
nssmat,
factor,
False,
np.squeeze(p.e[-1, :, :]),
np.squeeze(p.e[-1, :, :]).reshape((p.S, p.J)),
etr_params_3D,
mtrx_params_3D,
labor_noncompliance_rate_2D,
Expand All @@ -952,7 +952,7 @@ def SS_solver(
bssmat_s,
nssmat,
factor,
np.squeeze(p.e[-1, :, :]),
np.squeeze(p.e[-1, :, :]).reshape((p.S, p.J)),
etr_params_3D,
labor_noncompliance_rate_2D,
capital_noncompliance_rate_2D,
Expand All @@ -974,7 +974,7 @@ def SS_solver(
None,
False,
"SS",
np.squeeze(p.e[-1, :, :]),
np.squeeze(p.e[-1, :, :]).reshape((p.S, p.J)),
etr_params_3D,
p,
)
Expand Down Expand Up @@ -1020,7 +1020,7 @@ def SS_solver(
bqssmat,
rmssmat,
taxss,
np.squeeze(p.e[-1, :, :]),
np.squeeze(p.e[-1, :, :]).reshape((p.S, p.J)),
p.tau_c[-1, :],
p,
)
Expand Down Expand Up @@ -1077,7 +1077,7 @@ def SS_solver(
ubissmat,
theta,
etr_params_3D,
np.squeeze(p.e[-1, :, :]),
np.squeeze(p.e[-1, :, :]).reshape((p.S, p.J)),
p,
None,
"SS",
Expand Down
2 changes: 1 addition & 1 deletion ogcore/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,4 +20,4 @@
from ogcore.txfunc import *
from ogcore.utils import *

__version__ = "0.15.3"
__version__ = "0.15.4"
10 changes: 5 additions & 5 deletions ogcore/pensions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,14 +30,14 @@ def replacement_rate_vals(nssmat, wss, factor_ss, j, p):
"""
if j is not None:
e = np.squeeze(p.e[-1, :, j]) # Only computes using SS earnings
dim2 = 1
else:
e = np.squeeze(p.e[-1, :, :]) # Only computes using SS earnings
e = np.squeeze(p.e[-1, :, :]).reshape(
(p.S, p.J)
) # Only computes using SS earnings
dim2 = p.J
# adjust number of calendar years AIME computed from int model periods
equiv_periods = int(round((p.S / 80.0) * p.avg_earn_num_years)) - 1
if e.ndim == 2:
dim2 = e.shape[1]
else:
dim2 = 1
earnings = (e * (wss * nssmat * factor_ss)).reshape(p.S, dim2)
# get highest earning years for number of years AIME computed from
highest_earn = (
Expand Down
3 changes: 2 additions & 1 deletion ogcore/tax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -302,7 +302,7 @@ def net_taxes(
)
T_BQ = bequest_tax_liab(r, b, bq, t, j, method, p)
T_W = wealth_tax_liab(r, b, t, j, method, p)

# print("Net taxes shapes = ", T_I.shape, pension.shape, T_BQ.shape, T_W.shape)
net_tax = T_I - pension + T_BQ + T_W - tr - ubi

return net_tax
Expand Down Expand Up @@ -373,6 +373,7 @@ def income_tax_liab(r, w, b, n, factor, t, j, method, e, etr_params, p):
tax_filer = p.income_tax_filer[-1, :]
income = r * b + w * e * n
labor_income = w * e * n

T_I = (
ETR_income(
r,
Expand Down
7 changes: 5 additions & 2 deletions ogcore/txfunc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -263,7 +263,10 @@ def get_tax_rates(
txrates = tau_income + shift_income + shift
elif tax_func_type == "linear":
rate = np.squeeze(params[..., 0])
txrates = rate * np.ones_like(income)
try:
txrates = rate * np.ones_like(income)
except ValueError:
txrates = rate.reshape(income.shape) * np.ones_like(income)
elif tax_func_type == "mono":
if for_estimation:
mono_interp = params[0]
Expand Down Expand Up @@ -367,7 +370,7 @@ def get_tax_rates(
]
for t in range(income.shape[0])
]
txrates = np.squeeze(np.array(txrates))
txrates = np.squeeze(np.array(txrates)).reshape(X.shape)

return txrates

Expand Down
2 changes: 1 addition & 1 deletion setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

setuptools.setup(
name="ogcore",
version="0.15.3",
version="0.15.4",
author="Jason DeBacker and Richard W. Evans",
license="CC0 1.0 Universal (CC0 1.0) Public Domain Dedication",
description="A general equilibrium overlapping generations model for fiscal policy analysis",
Expand Down
18 changes: 18 additions & 0 deletions tests/test_SS.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -587,6 +587,21 @@ def test_solve_for_j():
"gamma_g": [0.0, 0.0, 0.0, 0.0],
}
filename8 = "inner_loop_outputs_reform_MneI.pkl"
param_updates9 = {
"J": 1,
"lambdas": np.array([1.0]),
"e": np.ones((80, 1)),
"beta_annual": [0.96],
"chi_b": [80],
"labor_income_tax_noncompliance_rate": [[0.0]],
"capital_income_tax_noncompliance_rate": [[0.0]],
"income_tax_filer": [[1]],
"wealth_tax_filer": [[1]],
"eta": np.ones((80, 1)) * (1 / 80),
"eta_RM": np.ones((80, 1)) * (1 / 80),
"replacement_rate_adjust": [[1.0]],
}
filename9 = "inner_loop_outputs_J1.pkl"


@pytest.mark.parametrize(
Expand All @@ -599,6 +614,7 @@ def test_solve_for_j():
(False, 0.04260341179572245, param_updates5, filename5),
(False, 0.04759112768438152, param_updates7, filename7),
(False, 0.04759112768438152, param_updates8, filename8),
(True, 0.04, param_updates9, filename9),
],
ids=[
"Baseline, Small Open",
Expand All @@ -608,6 +624,7 @@ def test_solve_for_j():
"Reform, baseline spending",
"Reform, M>1",
"Reform, I!=>M",
"J=1",
],
)
def test_inner_loop(baseline, r_p, param_updates, filename, dask_client):
Expand Down Expand Up @@ -662,6 +679,7 @@ def test_inner_loop(baseline, r_p, param_updates, filename, dask_client):
)
test_tuple = SS.inner_loop(outer_loop_vars, p, dask_client)

# The output format for the inner loop
try:
(
euler_errors,
Expand Down
46 changes: 32 additions & 14 deletions tests/test_TPI.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,8 +177,8 @@ def dask_client():
client = Client(cluster)
yield client
# teardown
client.close()
cluster.close()
# client.close()
# cluster.close()


filename1 = "intial_SS_values_baseline.pkl"
Expand Down Expand Up @@ -870,6 +870,21 @@ def test_run_TPI(baseline, param_updates, filename, tmpdir, dask_client):
filename8 = os.path.join(
CUR_PATH, "test_io_data", "run_TPI_outputs_baseline_Kg_nonzero_2.pkl"
)
param_updates10 = {
"J": 1,
"lambdas": np.array([1.0]),
"e": np.ones((40, 1)),
"beta_annual": [0.96],
"chi_b": [80],
"labor_income_tax_noncompliance_rate": [[0.0]],
"capital_income_tax_noncompliance_rate": [[0.0]],
"income_tax_filer": [[1]],
"wealth_tax_filer": [[1]],
"eta": np.ones((40, 1)) * (1 / 40),
"eta_RM": np.ones((40, 1)) * (1 / 40),
"replacement_rate_adjust": [[1.0]],
}
filename10 = os.path.join(CUR_PATH, "test_io_data", "run_TPI_outputs_J1.pkl")
# read in mono tax funcs (not age specific)
if sys.version_info[1] < 11:
dict_params = utils.safe_read_pickle(
Expand Down Expand Up @@ -899,6 +914,7 @@ def test_run_TPI(baseline, param_updates, filename, tmpdir, dask_client):
CUR_PATH, "test_io_data", "run_TPI_outputs_mono_2.pkl"
)


if sys.version_info[1] < 11:
test_list = [
(True, param_updates2, filename2),
Expand All @@ -922,22 +938,24 @@ def test_run_TPI(baseline, param_updates, filename, tmpdir, dask_client):
]
else:
test_list = [
# (True, param_updates2, filename2),
# (True, param_updates5, filename5),
# (True, param_updates6, filename6),
# (True, param_updates7, filename7),
# (True, {}, filename1),
# (False, param_updates4, filename4),
(True, param_updates2, filename2),
(True, param_updates5, filename5),
(True, param_updates6, filename6),
(True, param_updates7, filename7),
(True, {}, filename1),
(False, param_updates4, filename4),
(True, param_updates8, filename8),
(True, param_updates10, filename10),
]
id_list = [
# "Baseline, balanced budget",
# "Baseline, small open",
# "Baseline, small open for some periods",
# "Baseline, delta_tau = 0",
# "Baseline",
# "Reform, baseline spending",
"Baseline, balanced budget",
"Baseline, small open",
"Baseline, small open for some periods",
"Baseline, delta_tau = 0",
"Baseline",
"Reform, baseline spending",
"Baseline, Kg>0",
"J=1",
]


Expand Down
Binary file added tests/test_io_data/inner_loop_outputs_J1.pkl
View file
Open in desktop
Binary file not shown.
Binary file added tests/test_io_data/run_TPI_outputs_J1.pkl
View file
Open in desktop
Binary file not shown.
24 changes: 21 additions & 3 deletions tests/test_pensions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,27 @@
}
p.update_specifications(new_param_values)
p.retire = [3, 3, 3, 3, 3, 3, 3, 3]
p1 = copy.deepcopy(p)
p2 = copy.deepcopy(p)
p3 = copy.deepcopy(p)
p1 = Specifications()
p1.update_specifications(
{
"S": 4,
"rho": rho_vec.tolist(),
"lambdas": [0.5, 0.5],
"labor_income_tax_noncompliance_rate": [[0.0]],
"capital_income_tax_noncompliance_rate": [[0.0]],
"income_tax_filer": [[1]],
"wealth_tax_filer": [[1]],
"replacement_rate_adjust": [[1.0]],
"J": 2,
"T": 4,
"chi_n": np.ones(4),
"eta": (np.ones((4, 2)) / (4 * 2)),
"e": np.ones((4, 2)),
}
)
p1.retire = [3, 3, 3, 3, 3, 3, 3, 3]
p2 = copy.deepcopy(p1)
p3 = copy.deepcopy(p1)
# Use just a column of e
p1.e = np.transpose(np.array([[0.1, 0.3, 0.5, 0.2], [0.1, 0.3, 0.5, 0.2]]))
# e has two dimensions
Expand Down
Loading