Extrapolate benefits

taxcalc/records.py

@@ -113,7 +113,8 @@ def __init__(self,
                  gfactors=Growfactors(),
                  weights=PUF_WEIGHTS_FILENAME,
                  adjust_ratios=PUF_RATIOS_FILENAME,
-                 start_year=PUFCSV_YEAR):
+                 start_year=PUFCSV_YEAR,
+                 benefits=None):
         # pylint: disable=too-many-arguments
         self._data_year = start_year
         # read specified data
@@ -170,10 +171,17 @@ def __init__(self,
         if wt_colname in self.WT.columns:
             self.s006 = self.WT[wt_colname] * 0.01
 
+        # read in benefits data
+        # we would need to assert cps is being used and assert that the
+        # indices line up
+        self.benefits = benefits
+        self._read_benefits(benefits)
+
     @staticmethod
     def cps_constructor(data=None,
                         exact_calculations=False,
-                        growfactors=Growfactors()):
+                        growfactors=Growfactors(),
+                        benefits=None):
         """
         Static method returns a Records object instantiated with CPS
         input data.  This works in a analogous way to Records(), which
@@ -191,7 +199,8 @@ def cps_constructor(data=None,
                        gfactors=growfactors,
                        weights=Records.CPS_WEIGHTS_FILENAME,
                        adjust_ratios=Records.CPS_RATIOS_FILENAME,
-                       start_year=CPSCSV_YEAR)
+                       start_year=CPSCSV_YEAR,
+                       benefits=benefits)
 
     @property
     def data_year(self):
@@ -224,6 +233,20 @@ def increment_year(self):
             if wt_colname in self.WT.columns:
                 self.s006 = self.WT[wt_colname] * 0.01
 
+        # set ssi participation and benefits data for current_year
+        # in the future we would set corresponding attributes for
+        # SNAP, SS, Medicare, and Medicaid
+        # minimum year allowed is 2014
+        if self.benefits is not None:
+            year = max(2014, self.current_year)
+            recip_col_name = "ssi_recipients_{}".format(year)
+            self.ssi_recipients = (self.benefits.loc[:, recip_col_name]
+                                   .astype(np.float64).values)
+            ben_col_name = "ssi_benefits_{}".format(year)
+            self.ssi_benefits = (self.benefits.loc[:, ben_col_name]
+                                 .astype(np.float64).values)
+
+
     def set_current_year(self, new_current_year):
         """
         Set current year to specified value and updates FLPDYR variable.
@@ -514,3 +537,41 @@ def _read_ratios(self, ratios):
         if ADJ.index.name != 'agi_bin':
             ADJ.index.name = 'agi_bin'
         self.ADJ = ADJ
+
+    def _read_benefits(self, data):
+        """
+        Read benefits data from file or from a specified DataFrame or as None
+        if data is None
+        """
+        if data is None:
+            benefits = None
+            setattr(self, 'benefits', benefits)
+            return
+        if isinstance(data, pd.DataFrame):
+            benefits = data
+        elif isinstance(data, six.string_types):
+            data_path = os.path.join(Records.CUR_PATH, data)
+            if os.path.isfile(data_path):
+                # pylint: disable=redefined-variable-type
+                # (above because pylint mistakenly thinks ADJ not a DataFrame)
+                benefits = pd.read_csv(data_path)
+            else:
+                # cannot call read_egg_ function in unit tests
+                benefits = read_egg_csv(os.path.basename(data_path))
+        else:
+            msg = ('benefits is not None or a string'
+                   'or a Pandas DataFrame')
+            raise ValueError(msg)
+        assert isinstance(benefits, pd.DataFrame)
+        self.benefits = benefits
+        # set ssi participation and benefits data for current_year
+        # in the future we would set corresponding attributes for
+        # SNAP, SS, Medicare, and Medicaid
+        # minimum year allowed is 2014
+        year = max(2014, self.current_year)
+        recip_col_name = "ssi_recipients_{}".format(year)
+        self.ssi_recipients = (self.benefits.loc[:, recip_col_name]
+                               .astype(np.float64).values)
+        benefit_col_name = "ssi_benefits_{}".format(year)
+        self.ssi_benefits = (self.benefits.loc[:, benefit_col_name]
+                             .astype(np.float64).values)

taxcalc/tests/test_records.py

@@ -176,3 +176,34 @@ def test_csv_input_vars_md_contents(tests_path):
         for var in valid_less_civ:
             msg += 'VARIABLE= {}\n'.format(var)
         raise ValueError(msg)
+
+
+def test_benefits():
+    """
+    Ensure that benefits data is updated as expected
+    """
+    benefits_path = os.path.join(os.path.join(Records.CUR_PATH,
+                                 "cps_benefits.csv.gz"))
+    benefits = pd.read_csv(benefits_path, compression="gzip")
+    print(benefits.columns.tolist())
+    rec = Records.cps_constructor(data=None,
+                                  exact_calculations=False,
+                                  growfactors=Growfactors(),
+                                  benefits=benefits.copy(deep=True))
+
+    assert rec.current_year == 2014
+    recip_col_name = "ssi_recipients_{}".format(2014)
+    ben_col_name = "ssi_benefits_{}".format(2014)
+    assert (np.allclose(rec.ssi_recipients,
+                        benefits[recip_col_name].values) and
+            np.allclose(rec.ssi_benefits, benefits[ben_col_name].values))
+
+    for i in range(1, 13):
+        rec.increment_year()
+        recip_col_name = "ssi_recipients_{}".format(2014 + i)
+        ben_col_name = "ssi_benefits_{}".format(2014 + i)
+        assert (np.allclose(rec.ssi_recipients,
+                            benefits[recip_col_name].values) and
+                np.allclose(rec.ssi_benefits, benefits[ben_col_name].values))
+
+    assert rec.current_year == 2026