WIP: Add code to run OG-USA

environment.yml

@@ -8,9 +8,12 @@ dependencies:
 - behresp>=0.11.0
 - pandas>=0.23
 - numpy>=1.14
+- scipy
 - paramtools>=0.10.1
 - pytest
 - dask
+- distributed
+- multiprocessing
 - bokeh
 - markdown
 - tabulate

pytest.ini

@@ -0,0 +1,5 @@
+[pytest]
+testpaths =
+    taxbrain
+markers =
+    local: marks tests that run locally and not on GH Actions (mostly due to run time)

taxbrain/cli.py

@@ -47,7 +47,7 @@ def make_tables(tb, year, outpath):
 
 
 def cli_core(startyear, endyear, data, usecps, reform, behavior, assump,
-             baseline, outdir, name, make_report, author):
+             baseline, outdir, name, ogusa, make_report, author):
     """
     Core logic for the CLI
 
@@ -81,7 +81,7 @@ def cli_core(startyear, endyear, data, usecps, reform, behavior, assump,
     tb = TaxBrain(
         start_year=startyear, end_year=endyear, microdata=data,
         use_cps=usecps, reform=reform, behavior=behavior,
-        assump=assump, base_policy=baseline, verbose=True
+        assump=assump, ogusa=ogusa, base_policy=baseline, verbose=True
     )
     tb.run()
 
@@ -206,6 +206,13 @@ def cli_main():
         ),
         default=None
     )
+    parser.add_argument(
+        "--ogusa",
+        help=(
+            "If this argument is present, the model will be run using OG-USA"
+        ),
+        action="store_true"
+    )
     parser.add_argument(
         "--report",
         help=(
@@ -227,7 +234,7 @@ def cli_main():
     cli_core(
         args.startyear, args.endyear, args.data, args.usecps, args.reform,
         args.behavior, args.assump, args.baseline, args.outdir, args.name,
-        args.report
+        args.ogusa, args.author, args.report
     )
 
 

taxbrain/report.py

@@ -1,14 +1,23 @@
 import shutil
+from taxbrain.utils import is_paramtools_format
 import behresp
 import taxbrain
 import taxcalc as tc
 from pathlib import Path
-from .report_utils import (form_intro, form_baseline_intro, write_text, date,
-                           largest_tax_change, notable_changes,
-                           behavioral_assumptions, consumption_assumptions,
-                           policy_table, convert_table, growth_assumptions,
-                           md_to_pdf, DIFF_TABLE_ROW_NAMES,
-                           dollar_str_formatting)
+# from taxbrain.report_utils import (
+#     form_intro, form_baseline_intro, write_text, date,
+#     largest_tax_change, notable_changes,
+#     behavioral_assumptions, consumption_assumptions,
+#     policy_table, convert_table, growth_assumptions,
+#     md_to_pdf, DIFF_TABLE_ROW_NAMES,
+#     dollar_str_formatting)
+from taxbrain.report_utils import (
+    form_intro, form_baseline_intro, write_text, date,
+    largest_tax_change, notable_changes,
+    behavioral_assumptions, consumption_assumptions,
+    policy_table, convert_table, growth_assumptions,
+    DIFF_TABLE_ROW_NAMES,
+    dollar_str_formatting)
 
 
 CUR_PATH = Path(__file__).resolve().parent
@@ -42,7 +51,7 @@ def report(tb, name=None, change_threshold=0.05, description=None,
         created
     clean: bool
         boolean indicating whether all of the files written to create the
-        report should be deleated and a byte representation of the PDF returned
+        report should be deleted and a byte representation of the PDF returned
 
     Returns
     --------
@@ -143,7 +152,7 @@ def export_plot(plot, graph):
     pol_areas = list(pol_areas)
     # add policy areas to the intro text
     text_args["introduction"] = form_intro(pol_areas, description)
-    # write final sentance of introduction
+    # write final sentence of introduction
     current_law = tb.params["base_policy"]
     text_args["baseline_intro"] = form_baseline_intro(current_law)
 
@@ -230,7 +239,7 @@ def export_plot(plot, graph):
     if verbose:
         print("Compiling assumptions")
     text_args["behavior_assumps"] = behavioral_assumptions(tb)
-    # consumption asssumptions
+    # consumption assumptions
     text_args["consump_assumps"] = consumption_assumptions(tb)
     # growth assumptions
     text_args["growth_assumps"] = growth_assumptions(tb)

taxbrain/report_utils.py

@@ -2,7 +2,7 @@
 Helper Functions for creating the automated reports
 """
 import json
-import pypandoc
+# import pypandoc
 import numpy as np
 import pandas as pd
 import taxcalc as tc
@@ -11,7 +11,7 @@
 from datetime import datetime
 from tabulate import tabulate
 from collections import defaultdict, deque
-from .utils import is_paramtools_format
+from taxbrain import utils
 from typing import Union
 
 
@@ -43,30 +43,31 @@
                         '$100-200K', '$200-500K', '$500K-1M', '>$1M', 'ALL']
 
 
-def md_to_pdf(md_text, outputfile_path):
-    """
-    Convert Markdown version of report to a PDF. Returns bytes that can be
-    saved as a PDF
+# def md_to_pdf(md_text, outputfile_path):
+#     """
+#     Convert Markdown version of report to a PDF. Returns bytes that can be
+#     saved as a PDF
 
-    Parameters
-    ----------
-    md_text: str
-        report template written in markdown
-    outputfile_path: str
-        path to where the final file sohould be written
+#     Parameters
+#     ----------
+#     md_text: str
+#         report template written in markdown
+#     outputfile_path: str
+#         path to where the final file sohould be written
+
+#     Returns
+#     -------
+#     None
+#         Markdown text is saved as a PDF and the HTML used to create
+#         the report
+#     """
+#     # convert markdown text to pdf with pandoc
+#     pypandoc.convert_text(
+#         md_text, 'pdf', format='md', outputfile=outputfile_path,
+#         extra_args=['-V', 'geometry:margin=1.5cm',
+#                     '--pdf-engine', 'pdflatex']
+#     )
 
-    Returns
-    -------
-    None
-        Markdown text is saved as a PDF and the HTML used to create
-        the report
-    """
-    # convert markdown text to pdf with pandoc
-    pypandoc.convert_text(
-        md_text, 'pdf', format='md', outputfile=outputfile_path,
-        extra_args=['-V', 'geometry:margin=1.5cm',
-                    '--pdf-engine', 'pdflatex']
-    )
 
 
 def convert_table(df, tablefmt: str = "pipe") -> str:
@@ -126,7 +127,7 @@ def policy_table(params):
     }
     reform_years = set()
     reform_by_year = defaultdict(lambda: deque())
-    if is_paramtools_format(params):
+    if utils.is_paramtools_format(params):
         params = convert_params(params)
     pol = tc.Policy()  # policy object used for getting original value
     # loop through all of the policy parameters in a given reform

taxbrain/taxbrain.py

@@ -6,6 +6,7 @@
 from dask import compute, delayed
 from collections import defaultdict
 from taxbrain.utils import weighted_sum, update_policy
+from taxbrain.tbogusa import run_ogusa
 from typing import Union
 
 
@@ -23,10 +24,11 @@ class TaxBrain:
     }
 
     def __init__(self, start_year: int, end_year: int = LAST_BUDGET_YEAR,
-                 microdata: Union[str, dict] = None, use_cps: bool = False,
+                 microdata: Union[str, pd.DataFrame] = None,
+                 use_cps: bool = False,
                  reform: Union[str, dict] = None, behavior: dict = None,
-                 assump=None, base_policy: Union[str, dict] = None,
-                 verbose=False):
+                 assump=None, ogusa: bool = False,
+                 base_policy: Union[str, dict] = None, verbose=False):
         """
         Constructor for the TaxBrain class
 
@@ -56,6 +58,9 @@ def __init__(self, start_year: int, end_year: int = LAST_BUDGET_YEAR,
         assump: str
             A string pointing to a JSON file containing user specified
             economic assumptions.
+        ogusa: bool
+            A boolean value to indicate whether or not the analysis should
+            be run with OG-USA
         base_policy: str or dict
             Individual income tax policy to use as the baseline for
             the analysis. This policy will be implemented in the base
@@ -99,6 +104,7 @@ def __init__(self, start_year: int, end_year: int = LAST_BUDGET_YEAR,
         # Process user inputs early to throw any errors quickly
         self.params = self._process_user_mods(reform, assump)
         self.params["behavior"] = behavior
+        self.ogusa = ogusa
         if base_policy:
             base_policy = self._process_user_mods(base_policy, None)
             self.params["base_policy"] = base_policy["policy"]
@@ -107,7 +113,8 @@ def __init__(self, start_year: int, end_year: int = LAST_BUDGET_YEAR,
 
         self.has_run = False
 
-    def run(self, varlist: list = DEFAULT_VARIABLES):
+    def run(self, varlist: list = DEFAULT_VARIABLES, client=None,
+            num_workers=1):
         """
         Run the calculators. TaxBrain will determine whether to do a static or
         partial equilibrium run based on the user's inputs when initializing
@@ -122,6 +129,18 @@ def run(self, varlist: list = DEFAULT_VARIABLES):
         -------
         None
         """
+        if self.ogusa:
+            if self.verbose:
+                print("Running OG-USA")
+            if self.use_cps:
+                data_source = "cps"
+            else:
+                data_source = "puf"
+            og_results = run_ogusa(
+                iit_reform=self.params["policy"],
+                data_source=data_source, start_year=self.start_year,
+                client=client, num_workers=num_workers)
+            self._apply_ogusa(og_results)
         base_calc, reform_calc = self._make_calculators()
         if not isinstance(varlist, list):
             msg = f"'varlist' is of type {type(varlist)}. Must be a list."
@@ -450,3 +469,28 @@ def _make_calculators(self):
         # delete all unneeded variables
         del gd_base, gd_reform, records, gf_base, gf_reform, policy
         return base_calc, reform_calc
+
+    def _apply_ogusa(self, og_results):
+        """
+        Apply the results of the OG-USA run
+        Parameters
+        ----------
+        og_results: Numpy array
+            percentage changes in macro variables used to update grow
+            factors
+
+        Returns
+        -------
+        None
+        """
+        # changes in wage growth rates are at the 4th index
+        wage_change = og_results
+        gf = tc.GrowFactors()
+        grow_diff = {}
+        for i, yr in enumerate(range(self.start_year, self.end_year)):
+            cur_val = gf.factor_value("AWAGE", yr)
+            grow_diff[yr] = float(cur_val * (1 + wage_change[i]))
+        final_growdiffs = {
+            "AWAGE": grow_diff
+        }
+        self.params["growdiff_response"] = final_growdiffs

taxbrain/tbogusa.py

@@ -0,0 +1,95 @@
+import shutil
+import os
+from ogusa.execute import runner
+from ogusa.utils import safe_read_pickle
+from pathlib import Path
+
+
+CUR_PATH = Path(__file__).resolve().parent
+REFORM_DIR = Path(CUR_PATH, "ogusa_reform")
+BASE_DIR = Path(CUR_PATH, "ogusa_baseline")
+
+
+def run_ogusa(iit_base={}, iit_reform={}, og_spec_base={},
+              og_spec_reform={}, data_source='cps', start_year=2021,
+              num_years=10, client=None, num_workers=1):
+    """
+    Runs OG-USA model and returns percentage changes in macro variables
+    used as growth factors for microdata.
+
+    Parameters
+    ----------
+    iit_base: dict
+        baseline policy for Tax-Calculator
+    iit_reform: dict
+        reform policy for Tax-Calculator
+    og_spec_base: dict
+        OG-USA specifications for the baseline simulation
+    og_spec_reform: dict
+        OG-USA specifications for the reform simulation
+    data_source: str or Pandas DataFrame
+        path or DataFrame of microdata to be used by Tax-Calculator
+    start_year: int
+        year to start simulations in
+    num_years: int
+        number of years over which to return the results
+    client: Dask Client or None
+        Dask client to use
+    num_workers: int
+        number of workers for parallelization
+
+    Returns
+    -------
+    pct_w: Numpy array
+        percentage changes in wages, starting from start_year going to
+        start_year + num_years
+    """
+
+    '''
+    ------------------------------------------------------------------------
+    Run OG-USA baseline
+    ------------------------------------------------------------------------
+    '''
+    og_spec_base['start_year'] = start_year
+    og_spec_base['tax_func_type'] = 'GS'
+    og_spec_base['age_specific'] = False
+    kwargs = {'output_base': BASE_DIR, 'baseline_dir': BASE_DIR,
+              'test': False, 'time_path': True, 'baseline': True,
+              'og_spec': og_spec_base, 'guid': '',
+              'iit_reform': iit_base,
+              'run_micro': True, 'tax_func_path': None,
+              'data': data_source, 'client': client,
+              'num_workers': num_workers}
+    runner(**kwargs)
+
+    '''
+    ------------------------------------------------------------------------
+    Run reform policy
+    ------------------------------------------------------------------------
+    '''
+    og_spec_reform['start_year'] = start_year
+    og_spec_reform['tax_func_type'] = 'GS'
+    og_spec_reform['age_specific'] = False
+    kwargs = {'output_base': REFORM_DIR, 'baseline_dir': BASE_DIR,
+              'test': False, 'time_path': True, 'baseline': False,
+              'og_spec': og_spec_reform, 'guid': '',
+              'iit_reform': iit_reform, 'run_micro': False,
+              'tax_func_path': None, 'data': data_source,
+              'client': client, 'num_workers': num_workers}
+    runner(**kwargs)
+
+    # return ans - the percentage changes in macro aggregates and prices
+    # due to policy changes from the baseline to the reform
+    base_tpi = safe_read_pickle(
+        os.path.join(BASE_DIR, 'TPI', 'TPI_vars.pkl'))
+    reform_tpi = safe_read_pickle(
+        os.path.join(REFORM_DIR, 'TPI', 'TPI_vars.pkl'))
+
+    # compute pct change in wages over first num_years
+    pct_w = ((reform_tpi['w'] - base_tpi['w']) / base_tpi['w'])[:num_years]
+
+    # remove newly created directories
+    shutil.rmtree(BASE_DIR)
+    shutil.rmtree(REFORM_DIR)
+
+    return pct_w