Merge pull request #1113 from talumbau/consump_dropq

Consumption processing in dropq

taxcalc/dropq/dropq.py

@@ -1,5 +1,5 @@
 from __future__ import print_function
-from .. import (Calculator, Records, Policy, Behavior,
+from .. import (Calculator, Records, Policy, Behavior, Consumption,
                 TABLE_LABELS, TABLE_COLUMNS, STATS_COLUMNS,
                 DIFF_TABLE_LABELS)
 
@@ -11,6 +11,7 @@
 import hashlib
 import copy
 import time
+import collections
 from .dropq_utils import create_dropq_difference_table as dropq_diff_table
 from .dropq_utils import create_dropq_distribution_table as dropq_dist_table
 from .dropq_utils import *
@@ -41,6 +42,29 @@
 NUM_YEARS_DEFAULT = 1
 
 
+def call_over_iterable(fn):
+    """
+    A modifier for the only_* functions in this module. The idea is that
+    these functions may be passed a tuple of reform dictionaries, instead
+    of just a dictionary (since file-based reforms can result in a tuple
+    of reform dictionaries. In that case, this decorator will call its
+    wrapped function for each dictionary in the tuple, then collect
+    and return the results
+    """
+    def wrapper(user_mods, start_year):
+        if isinstance(user_mods, tuple):
+            ans = [fn(mod, start_year) for mod in user_mods]
+            params = ans[0]
+            for a in ans:
+                params.update(a)
+            return params
+        else:
+            return fn(user_mods, start_year)
+
+    return wrapper
+
+
+@call_over_iterable
 def only_growth_assumptions(user_mods, start_year):
     """
     Extract any reform parameters that are pertinent to growth
@@ -55,6 +79,7 @@ def only_growth_assumptions(user_mods, start_year):
     return ga
 
 
+@call_over_iterable
 def only_behavior_assumptions(user_mods, start_year):
     """
     Extract any reform parameters that are pertinent to behavior
@@ -69,6 +94,22 @@ def only_behavior_assumptions(user_mods, start_year):
     return ba
 
 
+@call_over_iterable
+def only_consumption_assumptions(user_mods, start_year):
+    """
+    Extract any reform parameters that are pertinent to behavior
+    assumptions
+    """
+    con_dd = Consumption.default_data(start_year=start_year)
+    ca = {}
+    for year, reforms in user_mods.items():
+        overlap = set(con_dd.keys()) & set(reforms.keys())
+        if overlap:
+            ca[year] = {param: reforms[param] for param in overlap}
+    return ca
+
+
+@call_over_iterable
 def only_reform_mods(user_mods, start_year):
     """
     Extract parameters that are just for policy reforms
@@ -77,25 +118,29 @@ def only_reform_mods(user_mods, start_year):
     beh_dd = Behavior.default_data(start_year=start_year)
     growth_dd = growth.Growth.default_data(start_year=start_year)
     policy_dd = policy.Policy.default_data(start_year=start_year)
+    param_code_names = policy.Policy.VALID_PARAM_CODE_NAMES
     for year, reforms in user_mods.items():
         all_cpis = {p for p in reforms.keys() if p.endswith("_cpi") and
                     p[:-4] in policy_dd.keys()}
         pols = set(reforms.keys()) - set(beh_dd.keys()) - set(growth_dd.keys())
-        pols &= set(policy_dd.keys())
+        pols &= set(policy_dd.keys()) | param_code_names
         pols ^= all_cpis
         if pols:
             pol_refs[year] = {param: reforms[param] for param in pols}
     return pol_refs
 
 
+@call_over_iterable
 def get_unknown_parameters(user_mods, start_year):
     """
-    Extract parameters that are just for policy reforms
+    Returns any parameters that are not known to Tax-Calculator
     """
     beh_dd = Behavior.default_data(start_year=start_year)
     growth_dd = growth.Growth.default_data(start_year=start_year)
     policy_dd = policy.Policy.default_data(start_year=start_year)
-    unknown_params = []
+    consump_dd = Consumption.default_data(start_year=start_year)
+    param_code_names = policy.Policy.VALID_PARAM_CODE_NAMES
+    unknown_params = collections.defaultdict(list)
     for year, reforms in user_mods.items():
         everything = set(reforms.keys())
         all_cpis = {p for p in reforms.keys() if p.endswith("_cpi")}
@@ -104,11 +149,12 @@ def get_unknown_parameters(user_mods, start_year):
         bad_cpis = all_cpis - all_good_cpis
         remaining = everything - all_cpis
         if bad_cpis:
-            unknown_params += list(bad_cpis)
+            unknown_params['bad_cpis'] += list(bad_cpis)
         pols = (remaining - set(beh_dd.keys()) - set(growth_dd.keys()) -
-                set(policy_dd.keys()))
+                set(policy_dd.keys()) - set(consump_dd.keys()) -
+                param_code_names)
         if pols:
-            unknown_params += list(pols)
+            unknown_params['policy'] += list(pols)
 
     return unknown_params
 
@@ -129,6 +175,23 @@ def elasticity_of_gdp_year_n(user_mods, year_n):
 
 
 def random_seed_from_plan(user_mods):
+    """
+    Handles the case of getting a tuple of reform mods
+    """
+
+    if isinstance(user_mods, tuple):
+        ans = 0
+        for mod in user_mods:
+            ans += _random_seed_from_plan(mod)
+        return ans % np.iinfo(np.uint32).max
+    else:
+        return _random_seed_from_plan(user_mods)
+
+
+def _random_seed_from_plan(user_mods):
+    """
+    Handles the case of getting a single reform mod dictionary
+    """
     all_vals = []
     for year in sorted(user_mods.keys()):
         all_vals.append(str(year))
@@ -445,6 +508,10 @@ def calculate_baseline_and_reform(year_n, start_year, is_strict,
     if growth_assumptions:
         calc1.growth.update_growth(growth_assumptions)
 
+    consump_assumptions = only_consumption_assumptions(user_mods, start_year)
+    if consump_assumptions:
+        calc1.consumption.update_consumption(consump_assumptions)
+
     while calc1.current_year < start_year:
         calc1.increment_year()
     calc1.calc_all()
@@ -456,6 +523,10 @@ def calculate_baseline_and_reform(year_n, start_year, is_strict,
     if growth_assumptions:
         calc2.growth.update_growth(growth_assumptions)
 
+    consump_assumptions = only_consumption_assumptions(user_mods, start_year)
+    if consump_assumptions:
+        calc2.consumption.update_consumption(consump_assumptions)
+
     while calc2.current_year < start_year:
         calc2.increment_year()
     calc2.calc_all()
@@ -480,6 +551,9 @@ def calculate_baseline_and_reform(year_n, start_year, is_strict,
     if growth_assumptions:
         calc3.growth.update_growth(growth_assumptions)
 
+    if consump_assumptions:
+        calc3.consumption.update_consumption(consump_assumptions)
+
     if behavior_assumptions:
         calc3.behavior.update_behavior(behavior_assumptions)
 

taxcalc/tests/test_dropq.py

@@ -1,5 +1,6 @@
 import os
 import json
+import tempfile
 import numpy as np
 import numpy.testing as npt
 import pandas as pd
@@ -11,6 +12,53 @@
 from taxcalc import multiyear_diagnostic_table
 
 
+REFORM_CONTENTS = """
+// Specify AGI exclusion of some fraction of investment income
+{
+  "policy": {
+    // specify fraction of investment income that can be excluded from AGI
+    "_ALD_Investment_ec_rt": {"2016": [0.50]},
+
+    // specify "investment income" base to mean the sum of three things:
+    // taxable interest (e00300), qualified dividends (e00650), and
+    // long-term capital gains (p23250) [see functions.py for other
+    // types of investment income that can be included in the base]
+    "param_code": {"ALD_Investment_ec_base_code":
+                   "e00300 + e00650 + p23250"},
+
+    "_ALD_Investment_ec_base_code_active": {"2016": [true]}
+    // the dollar exclusion is the product of the base defined by code
+    // and the fraction defined above
+  },
+  "behavior": {
+    "_BE_sub": {"2016": [0.25]}
+  },
+  "growth": {
+  },
+  "consumption": {
+    "_MPC_e20400": {"2016": [0.01]}
+  }
+}
+"""
+
+
+@pytest.yield_fixture
+def reform_file():
+    """
+    Temporary reform file for testing dropq functions with reform file.
+    """
+    rfile = tempfile.NamedTemporaryFile(mode='a', delete=False)
+    rfile.write(REFORM_CONTENTS)
+    rfile.close()
+    # must close and then yield for Windows platform
+    yield rfile
+    if os.path.isfile(rfile.name):
+        try:
+            os.remove(rfile.name)
+        except OSError:
+            pass  # sometimes we can't remove a generated temporary file
+
+
 @pytest.fixture(scope='session')
 def puf_path(tests_path):
     """
@@ -58,6 +106,28 @@ def test_run_dropq_nth_year(is_strict, rjson, growth_params,
         assert fiscal_tots is not None
 
 
+@pytest.mark.parametrize("is_strict", [True, False])
+def test_run_dropq_nth_year_from_file(is_strict, puf_1991_path, reform_file):
+
+    user_reform = Calculator.read_json_reform_file(reform_file.name)
+    user_mods = user_reform
+
+    # Create a Public Use File object
+    tax_data = pd.read_csv(puf_1991_path)
+    first = 2016
+    rjson = True
+
+    (_, _, _, _, _, _, _, _,
+     _, _, fiscal_tots) = dropq.run_models(tax_data,
+                                           start_year=first,
+                                           is_strict=is_strict,
+                                           user_mods=user_mods,
+                                           return_json=rjson,
+                                           num_years=3)
+
+    assert fiscal_tots is not None
+
+
 @pytest.mark.requires_pufcsv
 def test_full_dropq_puf(puf_path):
 
@@ -237,8 +307,11 @@ def test_unknown_parameters_with_cpi():
     first_year = 2013
     user_mods = {first_year: myvars}
     ans = get_unknown_parameters(user_mods, 2015)
+    final_ans = []
+    for a in ans.values():
+        final_ans += a
     exp = set(["NOGOOD_cpi", "NO", "ELASTICITY_GDP_WRT_AMTR"])
-    assert set(ans) == exp
+    assert set(final_ans) == exp
 
 
 def test_format_macro_results():