Merge pull request PSLmodels#1077 from jdebacker/reform_theta

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.0] - 2025-12-03 12:00:00
+
+### Added
+
+- A new parameter `baseline_theta` to the `Parameters` class that allows the user to specify whether to use the steady-state replacement rate parameters from the baseline solution in a reform run.  See PR [#1077](https://github.com/PSLmodels/OG-Core/pull/1077)
+
 ## [0.14.14] - 2025-11-24 12:00:00
 
 ### Added
@@ -481,6 +487,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.0]: https://github.com/PSLmodels/OG-Core/compare/v0.14.14...v0.15.0
 [0.14.14]: https://github.com/PSLmodels/OG-Core/compare/v0.14.13...v0.14.14
 [0.14.13]: https://github.com/PSLmodels/OG-Core/compare/v0.14.12...v0.14.13
 [0.14.12]: https://github.com/PSLmodels/OG-Core/compare/v0.14.11...v0.14.12

ogcore/SS.py

@@ -283,9 +283,27 @@ def inner_loop(outer_loop_vars, p, client):
     # from dask.base import dask_sizeof
 
     if client:
-        # Scatter p only once and only if client not equal None
+
+        # Before scattering, temporarily remove unpicklable schema objects
+        schema_backup = {}
+        for attr in ["_defaults_schema", "_validator_schema", "sel"]:
+            if hasattr(p, attr):
+                schema_backup[attr] = getattr(p, attr)
+                try:
+                    delattr(p, attr)
+                except:
+                    pass
+
+        # Scatter the parameters
         scattered_p_future = client.scatter(p, broadcast=True)
 
+        # Restore the schema objects (they're not needed by workers anyway)
+        for attr, value in schema_backup.items():
+            try:
+                setattr(p, attr, value)
+            except:
+                pass
+
         # Launch in parallel with submit (or map)
         futures = []
         for j in range(p.J):
@@ -1363,7 +1381,7 @@ def SS_initial_guesses(p, b_val=0.0055, n_val=0.4, r_tr_scalars=[1.0, 1.0]):
     if p.baseline:
         guesses.append(p.initial_guess_factor_SS)
 
-        return guesses, b_guess, n_guess
+    return guesses, b_guess, n_guess
 
 
 def run_SS(p, client=None):
@@ -1381,92 +1399,90 @@ def run_SS(p, client=None):
     """
     # For initial guesses of w, r, TR, and factor, we use values that
     # are close to some steady state values.
-    if p.baseline is False and p.reform_use_baseline_solution:
-        # Use the baseline solution to get starting values for the reform
+    # Use the baseline solution to get starting values for the reform
+    use_new_guesses = False  # initialize this flag, switches to true
+    # if baseline solution not work for reform
+    if p.baseline is False:
         baseline_ss_path = os.path.join(p.baseline_dir, "SS", "SS_vars.pkl")
         ss_solutions = utils.safe_read_pickle(baseline_ss_path)
-        if p.reform_use_baseline_solution:
-            # use baseline solution as starting values if dimensions match
-            try:
-                if ss_solutions["b_sp1"].shape == (
-                    p.S,
-                    p.J,
-                ) and np.squeeze(
-                    ss_solutions["Y_m"].shape
-                ) == (p.M):
-                    logging.info("Using previous solutions for SS")
-                    (
-                        b_guess,
-                        n_guess,
-                        r_p_guess,
-                        rguess,
-                        wguess,
-                        p_m_guess,
-                        BQguess,
-                        TRguess,
-                        Yguess,
-                        factor_ss,
-                    ) = (
-                        ss_solutions["b_sp1"],
-                        ss_solutions["n"],
-                        float(ss_solutions["r_p"]),
-                        float(ss_solutions["r"]),
-                        float(ss_solutions["w"]),
-                        ss_solutions["p_m"],
-                        ss_solutions["BQ"],
-                        float(ss_solutions["TR"]),
-                        float(ss_solutions["Y"]),
-                        ss_solutions["factor"],
-                    )
-                    use_new_guesses = False
-                    if p.baseline_spending:
-                        TR_baseline = TRguess
-                        Ig_baseline = ss_solutions["I_g"]
-                    else:
-                        TR_baseline = None
-                        Ig_baseline = None
-                    BQ_items = [BQguess] if p.use_zeta else list(BQguess)
-                    guesses = (
-                        [r_p_guess, rguess, wguess]
-                        + list(p_m_guess)
-                        + [Yguess]
-                        + BQ_items
-                        + [TRguess]
-                    )
-                    # Now solve for the steady state of the reform
-                    ss_params = (
-                        b_guess,
-                        n_guess,
-                        TR_baseline,
-                        Ig_baseline,
-                        factor_ss,
-                        p,
-                        client,
-                    )
-
-                    # Solve for steady state using root finder
-                    sol = opt.root(
-                        SS_fsolve,
-                        guesses,
-                        args=ss_params,
-                        method=p.SS_root_method,
-                        tol=p.mindist_SS,
-                    )
+    if p.baseline is False and p.reform_use_baseline_solution:
+        # use baseline solution as starting values if dimensions match
+        try:
+            if ss_solutions["b_sp1"].shape == (
+                p.S,
+                p.J,
+            ) and np.squeeze(
+                ss_solutions["Y_m"].shape
+            ) == (p.M):
+                logging.info("Using previous solutions for SS")
+                (
+                    b_guess,
+                    n_guess,
+                    r_p_guess,
+                    rguess,
+                    wguess,
+                    p_m_guess,
+                    BQguess,
+                    TRguess,
+                    Yguess,
+                    factor_ss,
+                ) = (
+                    ss_solutions["b_sp1"],
+                    ss_solutions["n"],
+                    float(ss_solutions["r_p"]),
+                    float(ss_solutions["r"]),
+                    float(ss_solutions["w"]),
+                    ss_solutions["p_m"],
+                    ss_solutions["BQ"],
+                    float(ss_solutions["TR"]),
+                    float(ss_solutions["Y"]),
+                    ss_solutions["factor"],
+                )
+                # set new attribute of parameters for SS theta
+                p.SS_theta = ss_solutions["theta"]
+                use_new_guesses = False
+                if p.baseline_spending:
+                    TR_baseline = TRguess
+                    Ig_baseline = ss_solutions["I_g"]
                 else:
-                    logging.warning(
-                        "Dimensions of previous solutions for SS do not match"
-                    )
-                    use_new_guesses = True
-            except KeyError:
+                    TR_baseline = None
+                    Ig_baseline = None
+                BQ_items = [BQguess] if p.use_zeta else list(BQguess)
+                guesses = (
+                    [r_p_guess, rguess, wguess]
+                    + list(p_m_guess)
+                    + [Yguess]
+                    + BQ_items
+                    + [TRguess]
+                )
+                # Now solve for the steady state of the reform
+                ss_params = (
+                    b_guess,
+                    n_guess,
+                    TR_baseline,
+                    Ig_baseline,
+                    factor_ss,
+                    p,
+                    client,
+                )
+
+                # Solve for steady state using root finder
+                sol = opt.root(
+                    SS_fsolve,
+                    guesses,
+                    args=ss_params,
+                    method=p.SS_root_method,
+                    tol=p.mindist_SS,
+                )
+            else:
                 logging.warning(
-                    "KeyError: previous solutions for SS not found"
+                    "Dimensions of previous solutions for SS do not match"
                 )
                 use_new_guesses = True
-    if (
-        p.baseline
-        or p.reform_use_baseline_solution is False
-        or use_new_guesses
-    ):
+        except KeyError:
+            logging.warning("KeyError: previous solutions for SS not found")
+            use_new_guesses = True
+    if p.baseline or not p.reform_use_baseline_solution or use_new_guesses:
         # Loop over initial guesses of r and TR until find a solution or until have
         # gone through all guesses. This should usually solve in the first guess
         SS_solved = False
@@ -1495,6 +1511,8 @@ def run_SS(p, client=None):
                     factor_ss = ss_solutions[
                         "factor"
                     ]  # don't guess factor, use baseline
+                    # set new attribute of parameters for SS theta
+                    p.SS_theta = ss_solutions["theta"]
                 if p.baseline_spending:
                     TR_baseline = ss_solutions["TR"]
                     Ig_baseline = ss_solutions["I_g"]

ogcore/TPI.py

@@ -744,9 +744,25 @@ def run_TPI(p, client=None):
     euler_errors = np.zeros((p.T, 2 * p.S, p.J))
     TPIdist_vec = np.zeros(p.maxiter)
 
-    # scatter parameters to workers
-    if client:
-        scattered_p_future = client.scatter(p, broadcast=True)
+    # Before scattering, temporarily remove unpicklable schema objects
+    schema_backup = {}
+    for attr in ["_defaults_schema", "_validator_schema", "sel"]:
+        if hasattr(p, attr):
+            schema_backup[attr] = getattr(p, attr)
+            try:
+                delattr(p, attr)
+            except:
+                pass
+
+    # Scatter the parameters
+    scattered_p_future = client.scatter(p, broadcast=True)
+
+    # Restore the schema objects (they're not needed by workers anyway)
+    for attr, value in schema_backup.items():
+        try:
+            setattr(p, attr, value)
+        except:
+            pass
 
     # TPI loop
     while (TPIiter < p.maxiter) and (TPIdist >= p.mindist_TPI):

ogcore/__init__.py

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

ogcore/default_parameters.json

@@ -333,7 +333,9 @@
                 "value": false
             }
         ],
-        "validators": {}
+        "validators": {
+           "choice": {"choices": [true, false]}
+            }
     },
     "beta_annual": {
         "title": "Annual rate of time preference for households",
@@ -3686,7 +3688,9 @@
                 "value": false
             }
         ],
-        "validators": {}
+        "validators": {
+           "choice": {"choices": [true, false]}
+            }
     },
     "constant_rates": {
         "title": "Flag to use linear tax functions",
@@ -3700,7 +3704,9 @@
                 "value": false
             }
         ],
-        "validators": {}
+        "validators": {
+           "choice": {"choices": [true, false]}
+            }
     },
     "zero_taxes": {
         "title": "Flag to run model without any individual income taxes",
@@ -3714,7 +3720,9 @@
                 "value": false
             }
         ],
-        "validators": {}
+        "validators": {
+           "choice": {"choices": [true, false]}
+            }
     },
     "analytical_mtrs": {
         "title": "Flag to use analytically derived marginal tax rates in tax functions",
@@ -3728,7 +3736,9 @@
                 "value": false
             }
         ],
-        "validators": {}
+        "validators": {
+           "choice": {"choices": [true, false]}
+            }
     },
     "age_specific": {
         "title": "Flag to use analytically derived marginal tax rates in tax functions",
@@ -3742,7 +3752,9 @@
                 "value": true
             }
         ],
-        "validators": {}
+        "validators": {
+           "choice": {"choices": [true, false]}
+            }
     },
     "retirement_age": {
         "title": "Age at which agents can collect Social Security benefits",
@@ -4119,6 +4131,22 @@
             }
         }
     },
+    "baseline_theta": {
+        "title": "Flag for use in reform simulations to keep Social Security system replacement rate constant between baseline and reform runs",
+        "description": "Flag for use in reform simulations to keep Social Security system replacement rate constant between baseline and reform runs.",
+        "section_1": "Fiscal Policy Parameters",
+        "section_2": "Government Pension Parameters",
+        "notes": "",
+        "type": "bool",
+        "value": [
+            {
+                "value": false
+            }
+        ],
+        "validators": {
+           "choice": {"choices": [true, false]}
+            }
+    },
     "start_year": {
         "title": "Calendar year in which to start model analysis",
         "description": "Calendar year in which to start model analysis.",
@@ -4395,7 +4423,9 @@
                 "value": true
             }
         ],
-        "validators": {}
+        "validators": {
+           "choice": {"choices": [true, false]}
+            }
     },
     "initial_guess_r_SS": {
         "title": "Initial guess of r for the SS solution",