Merge remote-tracking branch 'upstream/master' into param_plots

docs/book/content/api/utils.rst

@@ -13,9 +13,13 @@ ogcore.utils
 .. autoclass:: Inequality
   :members: gini, var_of_logs, ratio_pct1_pct2, pct, topshare
 
+.. autoclass:: CustomHttpAdapter
+  :members: init_poolmanager
+
 .. automodule:: ogcore.utils
   :members: mkdirs, pct_diff_func, convex_combo, read_file,
     pickle_file_compare, comp_array, comp_scalar, dict_compare,
     to_timepath_shape, get_initial_path, safe_read_pickle, rate_conversion,
     save_return_table, print_progress, fetch_files_from_web, not_connected,
-    avg_by_bin, extrapolate_arrays
+    avg_by_bin, extrapolate_arrays, get_legacy_session, shift_bio_clock,
+    pct_change_unstationarized

ogcore/utils.py

@@ -8,6 +8,7 @@
 from io import BytesIO, StringIO
 import numpy as np
 import pandas as pd
+from scipy.interpolate import CubicSpline
 import pickle
 from pkg_resources import resource_stream, Requirement
 import importlib.resources
@@ -1070,3 +1071,187 @@ def get_legacy_session():
     session = requests.session()
     session.mount("https://", CustomHttpAdapter(ctx))
     return session
+
+
+def shift_bio_clock(
+    param_in,
+    initial_effect_period=1,
+    final_effect_period=1,
+    total_effect=1,
+    min_age_effect_felt=20,
+    bound_below=False,
+    bound_above=False,
+    use_spline=False,
+):
+    """
+    This function takes an initial array of parameters that has a time
+    dimension and an age dimension. It then applies a shift along the
+    age dimension (i.e., a change in the "biological clock") and phases
+    this shift in over some period of time.
+
+    Args:
+        param_in (Numpy array): initial parameter values, first two
+            dimensions must be time then age, respectively
+        initial_effect_period (int): first model period when transition
+            to new parameter values occurs
+        final_effect_period (int): model period when effect is fully
+            phased in (so transition from param_in to param_out happens
+            between model periods `initial_effect_period` and
+            `final_effect_period`)
+        total_effect (float): total number of model periods to shift
+            the biological clock (allows for fractions of periods)
+        min_age_effect_felt (int): minimum age at which the effect of
+            the transition is felt in model periods
+        bound_below (bool): whether param_out bounded below by param_in
+        bound_above (bool): whether param_out bounded above by param_in
+        use_spline (bool): whether to use a cubic spline to interpolate
+            tail of param_out
+
+    Returns:
+        param_out (Numpy array): updated parameter values
+    """
+    assert (
+        total_effect >= 0
+    )  # this code would need to change to accommodate effects < 0
+    n_dims = param_in.ndim
+    assert n_dims >= 2
+    T = param_in.shape[1]
+    S = param_in.shape[1]
+
+    # create a linear transition path
+    t = (
+        final_effect_period - initial_effect_period + 1
+    )  # number of periods transition over
+    transition_path = np.linspace(0, 1.0, t)
+    print("Transition path = ", transition_path)
+    transition_arr = np.zeros_like(param_in, dtype=float)
+    for i in range(t):
+        print("TRANS = ", transition_path[i] * np.ones_like(param_in[0, ...]))
+        transition_arr[initial_effect_period + i, ...] = transition_path[
+            i
+        ] * np.ones_like(param_in[0, ...])
+    transition_arr[final_effect_period:, ...] = np.ones_like(
+        param_in[final_effect_period:, ...]
+    )
+
+    param_shift = param_in.copy()
+    # Accounting for shifts that are fractions of a model period
+    total_effect_ru = int(np.ceil(total_effect))
+    if total_effect > 0:
+        pct_effect = total_effect / total_effect_ru
+    else:
+        pct_effect = 0
+    print("Pct effect = ", pct_effect)
+    # apply the transition path to the initial parameters
+    # find diff from shifting bio clock back total_effect years
+    param_shift[:, min_age_effect_felt:, ...] = param_in[
+        :, (min_age_effect_felt - total_effect_ru) : S - total_effect_ru, ...
+    ]
+    if use_spline:
+        # use cubic spline to avoid plateau at end of lifecycle profile
+        T = param_in.shape[0]
+        if n_dims == 3:
+            J = param_in.shape[-1]
+            for k in range(initial_effect_period, T):
+                for j in range(J):
+                    spline = CubicSpline(
+                        np.arange(S - total_effect),
+                        param_shift[k, :-total_effect, j],
+                    )
+                    param_shift[k, -total_effect:, j] = spline(
+                        np.arange(S - total_effect, S)
+                    )
+        else:
+            for k in range(initial_effect_period, T):
+                spline = CubicSpline(
+                    np.arange(S - total_effect), param_shift[k, :-total_effect]
+                )
+                param_shift[k, -total_effect:] = spline(
+                    np.arange(S - total_effect, S)
+                )
+    # make sure values are not lower after shift
+    if bound_below:
+        param_shift = np.maximum(param_shift, param_in)
+    # make sure values are not higher after shift
+    if bound_above:
+        param_shift = np.minimum(param_shift, param_in)
+    # Now transition the shift over time using the transition path
+    param_out = param_in + (
+        transition_arr * pct_effect * (param_shift - param_in)
+    )
+
+    return param_out
+
+
+def pct_change_unstationarized(
+    tpi_base,
+    param_base,
+    tpi_reform,
+    param_reform,
+    output_vars=["K", "Y", "C", "L", "r", "w"],
+):
+    """
+    This function takes the time paths of variables from the baseline
+    and reform and parameters from the baseline and reform runs and
+    computes percent changes for each variable in the output_vars list.
+    The function first unstationarizes the time paths of the variables
+    and then computes the percent changes.
+
+    Args:
+        tpi_base (Numpy array): time path of the output variables from
+            the baseline run
+        param_base (Specifications object): dictionary of parameters
+            from the baseline run
+        tpi_reform (Numpy array): time path of the output variables from
+            the reform run
+        param_reform (Specifications object): dictionary of parameters
+            from the reform run
+        output_vars (list): list of variables for which to compute
+            percent changes
+
+    Returns:
+        pct_changes (dict): dictionary of percent changes for each
+            variable in output_vars list
+    """
+    # compute non-stationary variables
+    non_stationary_output = {"base": {}, "reform": {}}
+    pct_changes = {}
+    T = param_base.T
+    for var in output_vars:
+        if var in ["K", "C", "Y", "BQ", "TR", "UBI", "D", "total_tax_revenue"]:
+            non_stationary_output["base"][var] = (
+                tpi_base[var][:T]
+                * np.cumprod(1 + param_base.g_n[:T])
+                * np.exp(param_base.g_y * np.arange(param_base.T))
+            )
+            non_stationary_output["reform"][var] = (
+                tpi_reform[var][:T]
+                * np.cumprod(1 + param_reform.g_n[:T])
+                * np.exp(param_reform.g_y * np.arange(param_reform.T))
+            )
+        elif var in ["L"]:
+            non_stationary_output["base"][var] = tpi_base[var][
+                :T
+            ] * np.cumprod(1 + param_base.g_n[:T])
+            non_stationary_output["reform"][var] = tpi_reform[var][
+                :T
+            ] * np.cumprod(1 + param_reform.g_n[:T])
+        elif var in ["w", "ubi", "tr", "bq"]:
+            non_stationary_output["base"][var] = tpi_base[var][:T] * np.exp(
+                param_base.g_y * np.arange(param_base.T)
+            )
+            non_stationary_output["reform"][var] = tpi_reform[var][
+                :T
+            ] * np.exp(param_reform.g_y * np.arange(param_reform.T))
+        else:
+            non_stationary_output["base"][var] = tpi_base[var][:T]
+            non_stationary_output["reform"][var] = tpi_reform[var][:T]
+
+        # calculate percent change
+        pct_changes[var] = (
+            non_stationary_output["reform"][var]
+            / non_stationary_output["base"][var]
+            - 1
+        )
+
+    return pct_changes

tests/test_SS.py

@@ -199,45 +199,45 @@ def dask_client():
 args7 = (bssmat, nssmat, None, None, p7, None)
 expected7 = np.array(
     [
-        -0.07223291631036555,
-        -0.07826548741334771,
-        3.0911509552880334,
-        3.5079477421902254,
-        2.4414161493191413,
-        3.0137277396774236,
+        -0.06985935377445636,
+        -0.07388184648847439,
+        2.596215180212739,
+        3.0425352411195634,
+        2.08611520332783,
+        2.5993398246497392,
         0.0,
-        1.5448512748680931,
-        -0.0005707403143014808,
-        -0.01880468553296961,
-        -0.018630471665521498,
-        0.007811485269807893,
-        0.011911373110717808,
-        0.00898435233878905,
-        -0.0029658077084825795,
-        0.13903661473812837,
-        0.03462075279248166,
+        1.6276918281128583,
+        -0.0005336644680328222,
+        0.003641474531794135,
+        0.007892881165609,
+        0.007854285496066054,
+        0.011964025188377221,
+        0.00905400047723115,
+        0.0035962471039776792,
+        0.14649226453015723,
+        0.03816296076039217,
     ]
 )
 
 
 @pytest.mark.parametrize(
     "guesses,args,expected",
     [
-        # (guesses1, args1, expected1),
-        # (guesses2, args2, expected2),
-        # (guesses3, args3, expected3),
-        # (guesses4, args4, expected4),
-        # (guesses5, args5, expected5),
-        # (guesses6, args6, expected6),
+        (guesses1, args1, expected1),
+        (guesses2, args2, expected2),
+        (guesses3, args3, expected3),
+        (guesses4, args4, expected4),
+        (guesses5, args5, expected5),
+        (guesses6, args6, expected6),
         (guesses7, args7, expected7),
     ],
     ids=[
-        # "Baseline, Closed",
-        # "Reform, Closed",
-        # "Reform, Baseline spending=True, Closed",
-        # "Baseline, Partial Open",
-        # "Baseline, Small Open",
-        # "Baseline, Closed, delta_tau = 0",
+        "Baseline, Closed",
+        "Reform, Closed",
+        "Reform, Baseline spending=True, Closed",
+        "Baseline, Partial Open",
+        "Baseline, Small Open",
+        "Baseline, Closed, delta_tau = 0",
         "Baseline, M=4",
     ],
 )

tests/test_TPI.py

@@ -467,6 +467,7 @@ def test_run_TPI_full_run(
 )
 
 
+@pytest.mark.local
 @pytest.mark.parametrize(
     "baseline,param_updates,filename",
     [(True, {}, filename1), (False, {}, filename3)],