Merge pull request #1062 from jdebacker/J1

Merging

Author: rickecon

CHANGELOG.md

@@ -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
@@ -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

ogcore/SS.py

@@ -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,
     )
@@ -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,
     )
@@ -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,
     )
@@ -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,
     )
@@ -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",
@@ -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,
@@ -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,
@@ -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,
@@ -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,
     )
@@ -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,
     )
@@ -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",

ogcore/__init__.py

@@ -20,4 +20,4 @@
 from ogcore.txfunc import *
 from ogcore.utils import *
 
-__version__ = "0.15.3"
+__version__ = "0.15.4"

ogcore/pensions.py

@@ -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 = (

ogcore/tax.py

@@ -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
@@ -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,

ogcore/txfunc.py

@@ -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]
@@ -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
 

setup.py

@@ -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",

tests/test_SS.py

@@ -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(
@@ -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",
@@ -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):
@@ -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,

tests/test_TPI.py

@@ -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"
@@ -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(
@@ -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),
@@ -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",
     ]