add alt baseline example

examples/run_og_usa_current_policy_baseline.py

@@ -0,0 +1,181 @@
+import multiprocessing
+from distributed import Client
+import os
+import requests
+import json
+import time
+from taxcalc import Calculator
+from ogusa.calibrate import Calibration
+from ogcore.parameters import Specifications
+from ogcore import output_tables as ot
+from ogcore import output_plots as op
+from ogcore.execute import runner
+from ogcore.utils import safe_read_pickle
+
+
+def main():
+    # Define parameters to use for multiprocessing
+    num_workers = min(multiprocessing.cpu_count(), 7)
+    client = Client(n_workers=num_workers, threads_per_worker=1)
+    print("Number of workers = ", num_workers)
+
+    # Directories to save data
+    CUR_DIR = os.path.dirname(os.path.realpath(__file__))
+    base_dir = os.path.join(CUR_DIR, "OG-USA-CP-Example", "OUTPUT_BASELINE")
+    reform_dir = os.path.join(CUR_DIR, "OG-USA-CP-Example", "OUTPUT_REFORM")
+
+    """
+    ------------------------------------------------------------------------
+    Run baseline policy
+    ------------------------------------------------------------------------
+    """
+    # Set up baseline parameterization
+    p = Specifications(
+        baseline=True,
+        num_workers=num_workers,
+        baseline_dir=base_dir,
+        output_base=base_dir,
+    )
+    # Update parameters for baseline from default json file
+    p.update_specifications(
+        json.load(
+            open(
+                os.path.join(
+                    CUR_DIR, "..", "ogusa", "ogusa_default_parameters.json"
+                )
+            )
+        )
+    )
+    p.tax_func_type = "GS"
+    # get current policy JSON file
+    base_url = (
+        "github://PSLmodels:Tax-Calculator@master/taxcalc/"
+        + "reforms/ext.json"
+    )
+    ref = Calculator.read_json_param_objects(base_url, None)
+    iit_baseline = ref["policy"]
+    c = Calibration(
+        p,
+        estimate_tax_functions=True,
+        iit_baseline=iit_baseline,
+        client=client,
+    )
+    # close and delete client bc cache is too large
+    client.close()
+    del client
+    client = Client(n_workers=num_workers, threads_per_worker=1)
+    d = c.get_dict()
+    # # additional parameters to change
+    updated_params = {
+        "etr_params": d["etr_params"],
+        "mtrx_params": d["mtrx_params"],
+        "mtry_params": d["mtry_params"],
+        "mean_income_data": d["mean_income_data"],
+        "frac_tax_payroll": d["frac_tax_payroll"],
+    }
+    p.update_specifications(updated_params)
+    # Run model
+    start_time = time.time()
+    runner(p, time_path=True, client=client)
+    print("run time = ", time.time() - start_time)
+
+    """
+    ------------------------------------------------------------------------
+    Run reform policy
+    ------------------------------------------------------------------------
+    """
+    # Grab a reform JSON file already in Tax-Calculator
+    # In this example the 'reform' is a change to 2017 law
+    reform_url = (
+        "github://PSLmodels:examples@main/psl_examples/"
+        + "taxcalc/2017_law.json"
+    )
+    ref = Calculator.read_json_param_objects(reform_url, None)
+    iit_reform = ref["policy"]
+
+    # create new Specifications object for reform simulation
+    p2 = Specifications(
+        baseline=False,
+        num_workers=num_workers,
+        baseline_dir=base_dir,
+        output_base=reform_dir,
+    )
+    # Update parameters for baseline from default json file
+    p2.update_specifications(
+        json.load(
+            open(
+                os.path.join(
+                    CUR_DIR, "..", "ogusa", "ogusa_default_parameters.json"
+                )
+            )
+        )
+    )
+    p2.tax_func_type = "GS"
+    # Use calibration class to estimate reform tax functions from
+    # Tax-Calculator, specifying reform for Tax-Calculator in iit_reform
+    c2 = Calibration(
+        p2,
+        iit_baseline=iit_baseline,
+        iit_reform=iit_reform,
+        estimate_tax_functions=True,
+        client=client,
+    )
+    # close and delete client bc cache is too large
+    client.close()
+    del client
+    client = Client(n_workers=num_workers, threads_per_worker=1)
+    # update tax function parameters in Specifications Object
+    d = c2.get_dict()
+    # # additional parameters to change
+    updated_params = {
+        "cit_rate": [[0.35]],
+        "etr_params": d["etr_params"],
+        "mtrx_params": d["mtrx_params"],
+        "mtry_params": d["mtry_params"],
+        "mean_income_data": d["mean_income_data"],
+        "frac_tax_payroll": d["frac_tax_payroll"],
+    }
+    p2.update_specifications(updated_params)
+    # Run model
+    start_time = time.time()
+    runner(p2, time_path=True, client=client)
+    print("run time = ", time.time() - start_time)
+    client.close()
+
+    """
+    ------------------------------------------------------------------------
+    Save some results of simulations
+    ------------------------------------------------------------------------
+    """
+    base_tpi = safe_read_pickle(os.path.join(base_dir, "TPI", "TPI_vars.pkl"))
+    base_params = safe_read_pickle(os.path.join(base_dir, "model_params.pkl"))
+    reform_tpi = safe_read_pickle(
+        os.path.join(reform_dir, "TPI", "TPI_vars.pkl")
+    )
+    reform_params = safe_read_pickle(
+        os.path.join(reform_dir, "model_params.pkl")
+    )
+    ans = ot.macro_table(
+        base_tpi,
+        base_params,
+        reform_tpi=reform_tpi,
+        reform_params=reform_params,
+        var_list=["Y", "C", "K", "L", "r", "w"],
+        output_type="pct_diff",
+        num_years=10,
+        start_year=base_params.start_year,
+    )
+
+    # create plots of output
+    op.plot_all(
+        base_dir, reform_dir, os.path.join(CUR_DIR, "OG-USA_example_plots")
+    )
+
+    print("Percentage changes in aggregates:", ans)
+    # save percentage change output to csv file
+    ans.to_csv("ogusa_example_output.csv")
+
+
+if __name__ == "__main__":
+    # execute only if run as a script
+    main()