Merge pull request #565 from jdebacker/output_fixes

Merging.

.travis.yml

@@ -3,8 +3,8 @@ sudo: false
 dist: xenial
 language: python
 python:
-  - "3.6"
   - "3.7"
+  - "3.8"
 
 install:
   # Install conda

cs-config/cs_config/functions.py

@@ -65,7 +65,7 @@ class MetaParams(paramtools.Parameters):
 
 
 def get_version():
-    return "0.6.0"
+    return "0.6.2"
 
 
 def get_inputs(meta_param_dict):

ogusa/constants.py

@@ -33,6 +33,8 @@
               'w': 'Wage rate', 'BQ': 'Aggregate bequests ($BQ_{j,t}$)',
               'total_revenue': 'Total tax revenue ($REV_t$)',
               'business_revenue': 'Business tax revenue',
+              'iit_revenue': 'Individual income tax revenue',
+              'payroll_tax_revenue': 'Payroll tax revenue',
               'IITpayroll_revenue': 'IIT and payroll tax revenue',
               'n_mat': 'Labor Supply ($n_{j,s,t}$)',
               'c_path': 'Consumption ($c_{j,s,t}$)',

ogusa/fiscal.py

@@ -86,12 +86,14 @@ def D_G_path(r_gov, dg_fixed_values, p):
                                        G[t - 1] + TR[t - 1] -
                                        total_revenue[t - 1]))
             if (t >= p.tG1) and (t < p.tG2):
-                G[t] = (growth[t + 1] * (p.rho_G * p.debt_ratio_ss * Y[t] +
-                                         (1 - p.rho_G) * D[t]) -
-                        (1 + r_gov[t]) * D[t] + total_revenue[t] - TR[t])
+                G[t] = (
+                    growth[t + 1] * (p.rho_G * p.debt_ratio_ss * Y[t] +
+                                     (1 - p.rho_G) * D[t]) -
+                    (1 + r_gov[t]) * D[t] + total_revenue[t] - TR[t])
             elif t >= p.tG2:
-                G[t] = (growth[t + 1] * (p.debt_ratio_ss * Y[t]) -
-                        (1 + r_gov[t]) * D[t] + total_revenue[t] - TR[t])
+                G[t] = (
+                    growth[t + 1] * (p.debt_ratio_ss * Y[t]) -
+                    (1 + r_gov[t]) * D[t] + total_revenue[t] - TR[t])
             t += 1
 
         # in final period, growth rate has stabilized, so we can replace

ogusa/output_plots.py

@@ -335,7 +335,6 @@ def ability_bar(base_tpi, base_params, reform_tpi,
     plt.ylabel(r'Percentage Change in ' + VAR_LABELS[var])
     if plot_title:
         plt.title(plot_title, fontsize=15)
-    plt.legend(loc=9, bbox_to_anchor=(0.5, -0.15), ncol=2)
     if path:
         fig_path1 = os.path.join(path)
         plt.savefig(fig_path1, bbox_inches="tight")
@@ -656,15 +655,18 @@ def plot_all(base_output_path, reform_output_path, save_path):
 
     # Pct change in c, n, b, y, etr, mtrx, mtry by ability group over 10 years
     var_list = ['c_path', 'n_mat', 'bmat_splus1', 'etr_path',
-                'mtrx_path', 'mtry_path']
+                'mtrx_path', 'mtry_path', 'y_before_tax_mat']
     title_list = ['consumption', 'labor supply', 'savings',
                   'effective tax rates',
                   'marginal tax rates on labor income',
-                  'marginal tax rates on capital income']
-    path_list = ['Cons', 'Labor', 'Save', 'ETR', 'MTRx', 'MTRy']
+                  'marginal tax rates on capital income',
+                  'before tax income']
+    path_list = ['Cons', 'Labor', 'Save', 'ETR', 'MTRx', 'MTRy',
+                 'Income']
     for i, v in enumerate(var_list):
         ability_bar(base_tpi, base_params, reform_tpi, reform_params,
-                    var=v, num_years=10, start_year=DEFAULT_START_YEAR,
+                    var=v, num_years=10,
+                    start_year=base_params.start_year,
                     plot_title='Percentage changes in ' + title_list[i],
                     path=os.path.join(save_path, 'PctChange_' +
                                       path_list[i] + '.png'))

ogusa/output_tables.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pandas as pd
 import os
-from ogusa.constants import VAR_LABELS
+from ogusa.constants import VAR_LABELS, DEFAULT_START_YEAR
 from ogusa import wealth
 from ogusa.utils import save_return_table, Inequality
 cur_path = os.path.split(os.path.abspath(__file__))[0]
@@ -11,7 +11,7 @@ def macro_table(base_tpi, base_params, reform_tpi=None,
                 reform_params=None,
                 var_list=['Y', 'C', 'K', 'L', 'r', 'w'],
                 output_type='pct_diff', num_years=10, include_SS=True,
-                include_overall=True, start_year=2019,
+                include_overall=True, start_year=DEFAULT_START_YEAR,
                 table_format=None, path=None):
     '''
     Create a table of macro aggregates.
@@ -130,6 +130,8 @@ def macro_table(base_tpi, base_params, reform_tpi=None,
         # Make df with dict so can use pandas functions
         table_df = pd.DataFrame.from_dict(table_dict, orient='columns'
                                           ).set_index('Year').transpose()
+        table_df.reset_index(inplace=True)
+        table_df.rename(columns={'index': 'Variable'}, inplace=True)
         table = save_return_table(table_df, table_format, path)
 
     return table

ogusa/parameter_plots.py

@@ -4,12 +4,14 @@
 import matplotlib
 from ogusa.constants import GROUP_LABELS
 from ogusa import utils
+from ogusa.constants import DEFAULT_START_YEAR
 CUR_PATH = os.path.split(os.path.abspath(__file__))[0]
 style_file = os.path.join(CUR_PATH, 'OGUSAplots.mplstyle')
 plt.style.use(style_file)
 
 
-def plot_imm_rates(p, year=2019, include_title=False, path=None):
+def plot_imm_rates(p, year=DEFAULT_START_YEAR, include_title=False,
+                   path=None):
     '''
     Create a plot of immigration rates from OG-USA parameterization.
 
@@ -71,8 +73,9 @@ def plot_mort_rates(p, include_title=False, path=None):
         plt.savefig(fig_path)
 
 
-def plot_pop_growth(p, start_year=2019, num_years_to_plot=150,
-                    include_title=False, path=None):
+def plot_pop_growth(p, start_year=DEFAULT_START_YEAR,
+                    num_years_to_plot=150, include_title=False,
+                    path=None):
     '''
     Create a plot of population growth rates by year.
 

ogusa/parameter_tables.py

@@ -118,7 +118,7 @@ def tax_rate_table(base_TxFuncEst, base_params, reform_TxFuncEst=None,
         len_rates = len(base_rates[start_index: start_index +
                                    num_years])
         table = {
-                'Year': years[:len_rates],
+                 'Year': years[:len_rates],
                  'Baseline ' + VAR_LABELS[rate_type]:
                  base_rates[start_index: start_index + num_years],
                  'Reform ' + VAR_LABELS[rate_type]:
@@ -130,6 +130,8 @@ def tax_rate_table(base_TxFuncEst, base_params, reform_TxFuncEst=None,
     table_df.columns = table_df.iloc[0].astype('int').astype('str')
     table_df.reindex(table_df.index.drop('Year'))
     table_df.drop('Year', inplace=True)
+    table_df.reset_index(inplace=True)
+    table_df.rename(columns={'index': 'Variable'}, inplace=True)
     table_str = save_return_table(table_df, table_format, path,
                                   precision=2)
 

ogusa/tests/test_TPI.py

@@ -282,26 +282,15 @@ def test_run_TPI_full_run(baseline, param_updates, filename, tmp_path,
 param_updates4 = {'baseline_spending': True}
 filename4 = os.path.join(CUR_PATH, 'test_io_data',
                          'run_TPI_outputs_reform_baseline_spend_2.pkl')
-param_updates5 = {'zeta_K': [1.0]}
-filename5 = os.path.join(CUR_PATH, 'test_io_data',
-                         'run_TPI_outputs_baseline_small_open_2.pkl')
-param_updates6 = {'zeta_K': [0.2, 0.2, 0.2, 1.0, 1.0, 1.0, 0.2]}
-filename6 = os.path.join(
-    CUR_PATH, 'test_io_data',
-    'run_TPI_outputs_baseline_small_open_some_periods_2.pkl')
 
 
 @pytest.mark.parametrize('baseline,param_updates,filename',
                          [(True, param_updates2, filename2),
                           (True, param_updates1, filename1),
                           (False, param_updates3, filename3),
-                          (False, param_updates4, filename4),
-                          (True, param_updates5, filename5),
-                          (True, param_updates6, filename6)],
+                          (False, param_updates4, filename4)],
                          ids=['Baseline, balanced budget', 'Baseline',
-                              'Reform', 'Reform, baseline spending',
-                              'Baseline, small open',
-                              'Baseline, small open some periods'])
+                              'Reform', 'Reform, baseline spending'])
 def test_run_TPI(baseline, param_updates, filename, tmp_path,
                  dask_client):
     '''
@@ -349,3 +338,64 @@ def test_run_TPI(baseline, param_updates, filename, tmp_path,
         except ValueError:
             assert(np.allclose(test_dict[k], v[:p.T, :, :], rtol=1e-04,
                                atol=1e-04))
+
+
+param_updates5 = {'zeta_K': [1.0]}
+filename5 = os.path.join(CUR_PATH, 'test_io_data',
+                         'run_TPI_outputs_baseline_small_open_2.pkl')
+param_updates6 = {'zeta_K': [0.2, 0.2, 0.2, 1.0, 1.0, 1.0, 0.2]}
+filename6 = filename = os.path.join(
+    CUR_PATH, 'test_io_data',
+    'run_TPI_outputs_baseline_small_open_some_periods_2.pkl')
+
+
+@pytest.mark.full_run
+@pytest.mark.parametrize('baseline,param_updates,filename',
+                         [(True, param_updates5, filename5),
+                          (True, param_updates6, filename6)],
+                         ids=['Baseline, small open',
+                              'Baseline, small open for some periods'])
+def test_run_TPI_extra(baseline, param_updates, filename, tmp_path,
+                       dask_client):
+    '''
+    Test TPI.run_TPI function.  Provide inputs to function and
+    ensure that output returned matches what it has been before.
+    '''
+    baseline_dir = os.path.join(CUR_PATH, 'baseline')
+    output_base = baseline_dir
+    p = Specifications(baseline=True, baseline_dir=baseline_dir,
+                       output_base=output_base, test=True,
+                       client=dask_client, num_workers=NUM_WORKERS)
+    p.update_specifications(param_updates)
+    p.maxiter = 2  # this test runs through just two iterations
+    p.get_tax_function_parameters(
+        None, run_micro=False,
+        tax_func_path=os.path.join(
+            CUR_PATH, '..', 'data', 'tax_functions',
+            'TxFuncEst_baseline_CPS.pkl'))
+
+    # Need to run SS first to get results
+    SS.ENFORCE_SOLUTION_CHECKS = False
+    ss_outputs = SS.run_SS(p, None)
+
+    if p.baseline:
+        utils.mkdirs(os.path.join(p.baseline_dir, "SS"))
+        ss_dir = os.path.join(p.baseline_dir, "SS", "SS_vars.pkl")
+        with open(ss_dir, "wb") as f:
+            pickle.dump(ss_outputs, f)
+    else:
+        utils.mkdirs(os.path.join(p.output_base, "SS"))
+        ss_dir = os.path.join(p.output_base, "SS", "SS_vars.pkl")
+        with open(ss_dir, "wb") as f:
+            pickle.dump(ss_outputs, f)
+
+    TPI.ENFORCE_SOLUTION_CHECKS = False
+    test_dict = TPI.run_TPI(p, None)
+    expected_dict = utils.safe_read_pickle(filename)
+
+    for k, v in expected_dict.items():
+        try:
+            assert(np.allclose(test_dict[k], v, rtol=1e-04, atol=1e-04))
+        except ValueError:
+            assert(np.allclose(test_dict[k], v[:p.T, :, :], rtol=1e-04,
+                               atol=1e-04))

setup.py

@@ -4,9 +4,9 @@
     from distutils.core import setup
 
 with open('README.md') as f:
-        longdesc = f.read()
+    longdesc = f.read()
 
-version = '0.5.13'
+version = '0.6.2'
 
 config = {
     'description': 'General equilibribum, overlapping generations model for the USA',
@@ -36,8 +36,8 @@
         'Operating System :: OS Independent',
         'Programming Language :: Python',
         'Programming Language :: Python :: 3',
-        'Programming Language :: Python :: 3.6',
         'Programming Language :: Python :: 3.7',
+        'Programming Language :: Python :: 3.8',
         'Topic :: Software Development :: Libraries :: Python Modules'],
     'tests_require': ['pytest']
 }