<commit_before>import pandas as pd
import numpy as np
import os
from copy import deepcopy
import time

class Benefits():
    GROWTH_RATES_PATH = 'growth_rates.csv'
    CPS_BENEFIT_PATH = '../cps_data/cps_raw.csv.gz'
    CPS_WEIGHTS_PATH = '../cps_stage2/cps_weights.csv.gz'

    def __init__(self,
                 growth_rates=GROWTH_RATES_PATH,
                 cps_benefit=CPS_BENEFIT_PATH,
                 cps_weights=CPS_WEIGHTS_PATH,
                 benefit_names=["ssi", "medicaid", "medicare",
                                "vb", "snap", "ss"]):
        benefit_names = [bn.lower() for bn in benefit_names]
        self._read_data(growth_rates, cps_benefit, cps_weights,
                        benefit_names=benefit_names)
        self.current_year = 2014
        self.benefit_names = benefit_names


    def increment_year(self, tol=0.01):
        self.current_year += 1
        print("starting year", self.current_year)
        WT = self.WT.loc[:, 'WT'+str(self.current_year)] * 0.01
        for benefit in self.benefit_names:
            participant_targets = getattr(self, "{}_participant_targets".format(benefit))
            benefit_targets = getattr(self, "{}_benefit_targets".format(benefit))
            participation = getattr(self, "{}_participation".format(benefit))
            benefits = getattr(self, "{}_benefits".format(benefit))
            prob = getattr(self, "{}_prob".format(benefit))

            diff = participant_targets[self.current_year] - \
                (participation.sum(axis=1) * WT).sum()
            if abs(diff)/participant_targets[self.current_year] > tol:
                participation, benefits = \
                    self._extrapolate(WT,
                                      participation.copy(deep=True),
                                      benefits.copy(deep=True),
                                      prob,
                                      participant_targets[self.current_year],
                                      tol=tol)
            # part_wt = pd.concat(participation.sum(axis=1), WT)
            # extrapolated = part_wt.
            extrapolated = (participation.sum(axis=1) * WT).sum()
            print(benefit, 'this year total is ', extrapolated,
                  ', target is ', participant_targets[self.current_year],
                  ', diff is ',
                  (extrapolated - participant_targets[self.current_year]), ',',
                  ((extrapolated - participant_targets[self.current_year]) /
                   participant_targets[self.current_year]))

            lab = benefit + '_recipients_' + str(self.current_year)
            self.benefit_extrapolation[lab] = participation.sum(axis=1)

            total_current_benefits = (benefits.sum(axis=1) * WT).sum()
            lab = benefit + '_benefits_' + str(self.current_year)
            self.benefit_extrapolation[lab] = \
                (benefits.sum(axis=1) *
                 benefit_targets[self.current_year] / total_current_benefits)

            setattr(self, '{}_participation'.format(benefit), participation)
            setattr(self, '{}_benefits'.format(benefit), benefits)

    @staticmethod
    def _extrapolate(WT, I, benefits, prob, target, tol=0.01, J=15):
        """
        Goal: get number of participants as close to target as possible
        Steps:
            1. Data is currently in N x 15 shape. Reshape it to be
                N * 15 dim vector
            2. Determine if adding or removing
                if adding, only consider those who are not receiving benefits
                if removing, only consider those who are receiving benefits
            2. Sort in descending order of probability of receiving
                benefits in the future
            3. Calculate cumulative sum over all participants
            4. Find index where minimum absolute difference between
                cum. sum and target is minimized
            5. All those with index <= minimimizing index receive benefits
                and all those with index > minimizing index do not
            6. Stack dataframes, unravel (go back to N x 15 dim), return result

            Note: Unraveling is done by keeping track of each individual's
                location in the matrix (i, j).  Sorting by i and then j gives
                us a vector of tuples (1,1), (1,2),...,(1,15),(2,1),(2,2),
                ...,(N,15). Then by reshaping to n rows, 15 columns we
                get the same matrix that we started with but with the
                updated values
        """
        start = time.time()

        extrap_df = Benefits._stack_df(WT, I, benefits, prob, J)

        receives = extrap_df.loc[extrap_df.I > 0, ]
        avg_benefit = receives.benefits_wt.sum() / receives.I_wt.sum()

        actual = extrap_df.I_wt.sum()
        diff = actual - target
        if diff < 0:
            remove = False
            candidates = extrap_df.loc[extrap_df.I == 0, ]
            assert candidates.I.sum() == 0
            # set everyone as a participant and then remove surplus below
            # this is necessary since I * wt will all be zeroes
            candidates.loc[:, "I"] = np.ones(len(candidates))
            candidates.loc[:, "I_wt"] = candidates.I * candidates.wt
            noncandidates = extrap_df.loc[extrap_df.I == 1, ]
            assert noncandidates.I.sum() == len(noncandidates)
        else:
            remove = True
            candidates = extrap_df.loc[extrap_df.I == 1, ]
            assert candidates.I.sum() == len(candidates)
            noncandidates = extrap_df.loc[extrap_df.I == 0, ]
            assert noncandidates.I.sum() == 0
        noncan_part = noncandidates.I_wt.sum()
        del extrap_df
        # sort by probability of getting benefits in descending order
        # of prob so that we only keep those with highest prob of
        # receiving benefits in the future
        candidates.sort_values(by="prob", ascending=False, inplace=True)
        # create running sum of weighted participants
        candidates["cum_participants"] = candidates.I_wt.cumsum() + noncan_part
        # get absolute difference between candidates and target
        candidates["_diff"] = candidates.cum_participants - target

        # # check to make results are close enough
        assert np.allclose(
            candidates[candidates._diff <= 0].cum_participants.max(), target,
            atol=0.0, rtol=tol
        )

        # individuals prior to min_diff have highest probability of getting
        # benefits in the future ==> they get benefits
        candidates.loc[candidates._diff <= 0, "I"] = 1
        candidates.loc[candidates._diff > 0, "I"] = 0

        if remove:
            candidates.loc[candidates._diff > 0, 'benefits'] = 0<commit_msg>Copy dataframes<commit_after>import pandas as pd
import numpy as np
import os
from copy import deepcopy
import time

class Benefits():
    GROWTH_RATES_PATH = 'growth_rates.csv'
    CPS_BENEFIT_PATH = '../cps_data/cps_raw.csv.gz'
    CPS_WEIGHTS_PATH = '../cps_stage2/cps_weights.csv.gz'

    def __init__(self,
                 growth_rates=GROWTH_RATES_PATH,
                 cps_benefit=CPS_BENEFIT_PATH,
                 cps_weights=CPS_WEIGHTS_PATH,
                 benefit_names=["ssi", "medicaid", "medicare",
                                "vb", "snap", "ss"]):
        benefit_names = [bn.lower() for bn in benefit_names]
        self._read_data(growth_rates, cps_benefit, cps_weights,
                        benefit_names=benefit_names)
        self.current_year = 2014
        self.benefit_names = benefit_names


    def increment_year(self, tol=0.01):
        self.current_year += 1
        print("starting year", self.current_year)
        WT = self.WT.loc[:, 'WT'+str(self.current_year)] * 0.01
        for benefit in self.benefit_names:
            participant_targets = getattr(self, "{}_participant_targets".format(benefit))
            benefit_targets = getattr(self, "{}_benefit_targets".format(benefit))
            participation = getattr(self, "{}_participation".format(benefit))
            benefits = getattr(self, "{}_benefits".format(benefit))
            prob = getattr(self, "{}_prob".format(benefit))

            diff = participant_targets[self.current_year] - \
                (participation.sum(axis=1) * WT).sum()
            if abs(diff)/participant_targets[self.current_year] > tol:
                participation, benefits = \
                    self._extrapolate(WT,
                                      participation.copy(deep=True),
                                      benefits.copy(deep=True),
                                      prob,
                                      participant_targets[self.current_year],
                                      tol=tol)
            # part_wt = pd.concat(participation.sum(axis=1), WT)
            # extrapolated = part_wt.
            extrapolated = (participation.sum(axis=1) * WT).sum()
            print(benefit, 'this year total is ', extrapolated,
                  ', target is ', participant_targets[self.current_year],
                  ', diff is ',
                  (extrapolated - participant_targets[self.current_year]), ',',
                  ((extrapolated - participant_targets[self.current_year]) /
                   participant_targets[self.current_year]))

            lab = benefit + '_recipients_' + str(self.current_year)
            self.benefit_extrapolation[lab] = participation.sum(axis=1)

            total_current_benefits = (benefits.sum(axis=1) * WT).sum()
            lab = benefit + '_benefits_' + str(self.current_year)
            self.benefit_extrapolation[lab] = \
                (benefits.sum(axis=1) *
                 benefit_targets[self.current_year] / total_current_benefits)

            setattr(self, '{}_participation'.format(benefit), participation)
            setattr(self, '{}_benefits'.format(benefit), benefits)

    @staticmethod
    def _extrapolate(WT, I, benefits, prob, target, tol=0.01, J=15):
        """
        Goal: get number of participants as close to target as possible
        Steps:
            1. Data is currently in N x 15 shape. Reshape it to be
                N * 15 dim vector
            2. Determine if adding or removing
                if adding, only consider those who are not receiving benefits
                if removing, only consider those who are receiving benefits
            2. Sort in descending order of probability of receiving
                benefits in the future
            3. Calculate cumulative sum over all participants
            4. Find index where minimum absolute difference between
                cum. sum and target is minimized
            5. All those with index <= minimimizing index receive benefits
                and all those with index > minimizing index do not
            6. Stack dataframes, unravel (go back to N x 15 dim), return result

            Note: Unraveling is done by keeping track of each individual's
                location in the matrix (i, j).  Sorting by i and then j gives
                us a vector of tuples (1,1), (1,2),...,(1,15),(2,1),(2,2),
                ...,(N,15). Then by reshaping to n rows, 15 columns we
                get the same matrix that we started with but with the
                updated values
        """
        start = time.time()

        extrap_df = Benefits._stack_df(WT, I, benefits, prob, J)

        receives = extrap_df.loc[extrap_df.I > 0, ]
        avg_benefit = receives.benefits_wt.sum() / receives.I_wt.sum()

        actual = extrap_df.I_wt.sum()
        diff = actual - target
        if diff < 0:
            remove = False
            candidates = extrap_df.loc[extrap_df.I == 0, ].copy()
            assert candidates.I.sum() == 0
            # set everyone as a participant and then remove surplus below
            # this is necessary since I * wt will all be zeroes
            candidates.loc[:, "I"] = np.ones(len(candidates))
            candidates.loc[:, "I_wt"] = candidates.I * candidates.wt
            noncandidates = extrap_df.loc[extrap_df.I == 1, ].copy()
            assert noncandidates.I.sum() == len(noncandidates)
        else:
            remove = True
            candidates = extrap_df.loc[extrap_df.I == 1, ].copy()
            assert candidates.I.sum() == len(candidates)
            noncandidates = extrap_df.loc[extrap_df.I == 0, ].copy()
            assert noncandidates.I.sum() == 0
        noncan_part = noncandidates.I_wt.sum()
        del extrap_df
        # sort by probability of getting benefits in descending order
        # of prob so that we only keep those with highest prob of
        # receiving benefits in the future
        candidates.sort_values(by="prob", ascending=False, inplace=True)
        # create running sum of weighted participants
        candidates["cum_participants"] = candidates.I_wt.cumsum() + noncan_part
        # get absolute difference between candidates and target
        candidates["_diff"] = candidates.cum_participants - target

        # # check to make results are close enough
        assert np.allclose(
            candidates[candidates._diff <= 0].cum_participants.max(), target,
            atol=0.0, rtol=tol
        )

        # individuals prior to min_diff have highest probability of getting
        # benefits in the future ==> they get benefits
        candidates.loc[candidates._diff <= 0, "I"] = 1
        candidates.loc[candidates._diff > 0, "I"] = 0

        if remove:
            candidates.loc[candidates._diff > 0, 'benefits'] = 0
{'fraction_non_alphanumeric': 0.09608016504568229, 'fraction_numerical': 0.010020630710285883, 'mean_word_length': 1.9791483757682178, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1178130', 'n_tokens_mistral': 3989, 'n_tokens_neox': 3674, 'n_words': 998}