<commit_before>import census
import pandas as pd
import numpy as np
import us


# TODO DOCSTRING!!
class Census:

    def __init__(self, key, acsyear = 2016):
        self.c = census.Census(key)
        self.base_url = "https://s3-us-west-1.amazonaws.com/synthpop-data2/"
        self.pums_relationship_file_url = self.base_url + "tract10_to_puma.csv"
        self.pums_relationship_df = None
        self.pums10_population_base_url = \
            self.base_url + "puma10_p_%s_%s.csv"
        self.pums10_household_base_url = \
            self.base_url + "puma10_h_%s_%s.csv"
        self.pums00_population_base_url = \
            self.base_url + "puma00_p_%s_%s.csv"
        self.pums00_household_base_url = \
            self.base_url + "puma00_h_%s_%s.csv"
        self.pums_population_state_base_url = \
            self.base_url + "puma_p_%s.csv"
        self.pums_household_state_base_url = \
            self.base_url + "puma_h_%s.csv"
        self.fips_url = self.base_url + "national_county.txt"
        self.fips_df = None
        self.pums_cache = {}
        self.acsyear = acsyear

    # df1 is the disaggregate data frame (e.g. block groups)
    # df2 is the aggregate data frame (e.g. tracts)
    # need to scale down df2 variables to df1 level totals
    def _scale_and_merge(self, df1, tot1, df2, tot2, columns_to_scale,
                         merge_columns, suffixes):
        df = pd.merge(df1, df2, left_on=merge_columns, right_on=merge_columns,
                      suffixes=suffixes)

        # going to scale these too so store current values
        tot2, tot1 = df[tot2], df[tot1]
        # if agg number if 0, disaggregate should be 0
        # note this is filled by fillna below
        assert np.all(tot1[tot2 == 0] == 0)

        for col in columns_to_scale:
            df[col] = df[col] / tot2 * tot1
            # round?
            df[col] = df[col].fillna(0).astype('int')
        return df

    def block_group_query(self, census_columns, state, county, year= self.acsyear, 
                        tract=None,id=None):
        if id is None:
            id = "*"
        return self._query(census_columns, state, county,
                           forstr="block group:%s" % id,
                           tract=tract, year=year)

    def tract_query(self, census_columns, state, county, year= self.acsyear,
                    tract=None):
        if tract is None:
            tract = "*"
        return self._query(census_columns, state, county,
                           forstr="tract:%s" % tract,
                           year=year)

    def _query(self, census_columns, state, county, forstr,
               year, tract=None):
        c = self.c

        state, county = self.try_fips_lookup(state, county)

        if tract is None:
            in_str = 'state:%s county:%s' % (state, county)
        else:
            in_str = 'state:%s county:%s tract:%s' % (state, county, tract)

        dfs = []

        # unfortunately the api only queries 50 columns at a time
        # leave room for a few extra id columns
        def chunks(l, n):
            """ Yield successive n-sized chunks from l.
            """
            for i in range(0, len(l), n):
                yield l[i:i+n]

        for census_column_batch in chunks(census_columns, 45):
            census_column_batch = list(census_column_batch)
            d = c.acs5.get(['NAME'] + census_column_batch,
                          geo={'for': forstr,
                               'in': in_str}, year=year)
            df = pd.DataFrame(d)
            df[census_column_batch] = df[census_column_batch].astype('int')
            dfs.append(df)

        assert len(dfs) >= 1
        df = dfs[0]
        for mdf in dfs[1:]:
            df = pd.merge(df, mdf, on="NAME", suffixes=("", "_ignore"))
            drop_cols = list(filter(lambda x: "_ignore" in x, df.columns))
            df = df.drop(drop_cols, axis=1)

        return df

    def block_group_and_tract_query(self, block_group_columns,
                                    tract_columns, state, county,
                                    merge_columns, block_group_size_attr,
                                    tract_size_attr, year= self.acsyear, tract=None):
        df2 = self.tract_query(tract_columns, state, county, tract=tract,
                               year=year)
        df1 = self.block_group_query(block_group_columns, state, county,
                                     tract=tract, year=year)

        df = self._scale_and_merge(df1, block_group_size_attr, df2,
                                   tract_size_attr, tract_columns,
                                   merge_columns, suffixes=("", "_ignore"))
        drop_cols = list(filter(lambda x: "_ignore" in x, df.columns))
        df = df.drop(drop_cols, axis=1)

        return df

    def _get_pums_relationship(self):
        if self.pums_relationship_df is None:
            self.pums_relationship_df = \
                pd.read_csv(self.pums_relationship_file_url, dtype={
                    "statefp": "object",<commit_msg>unchanged censushelper, travis error<commit_after>import census
import pandas as pd
import numpy as np
import us


# TODO DOCSTRING!!
class Census:

    def __init__(self, key):
        self.c = census.Census(key)
        self.base_url = "https://s3-us-west-1.amazonaws.com/synthpop-data2/"
        self.pums_relationship_file_url = self.base_url + "tract10_to_puma.csv"
        self.pums_relationship_df = None
        self.pums10_population_base_url = \
            self.base_url + "puma10_p_%s_%s.csv"
        self.pums10_household_base_url = \
            self.base_url + "puma10_h_%s_%s.csv"
        self.pums00_population_base_url = \
            self.base_url + "puma00_p_%s_%s.csv"
        self.pums00_household_base_url = \
            self.base_url + "puma00_h_%s_%s.csv"
        self.pums_population_state_base_url = \
            self.base_url + "puma_p_%s.csv"
        self.pums_household_state_base_url = \
            self.base_url + "puma_h_%s.csv"
        self.fips_url = self.base_url + "national_county.txt"
        self.fips_df = None
        self.pums_cache = {}

    # df1 is the disaggregate data frame (e.g. block groups)
    # df2 is the aggregate data frame (e.g. tracts)
    # need to scale down df2 variables to df1 level totals
    def _scale_and_merge(self, df1, tot1, df2, tot2, columns_to_scale,
                         merge_columns, suffixes):
        df = pd.merge(df1, df2, left_on=merge_columns, right_on=merge_columns,
                      suffixes=suffixes)

        # going to scale these too so store current values
        tot2, tot1 = df[tot2], df[tot1]
        # if agg number if 0, disaggregate should be 0
        # note this is filled by fillna below
        assert np.all(tot1[tot2 == 0] == 0)

        for col in columns_to_scale:
            df[col] = df[col] / tot2 * tot1
            # round?
            df[col] = df[col].fillna(0).astype('int')
        return df

    def block_group_query(self, census_columns, state, county, year= 2016, 
                        tract=None,id=None):
        if id is None:
            id = "*"
        return self._query(census_columns, state, county,
                           forstr="block group:%s" % id,
                           tract=tract, year=year)

    def tract_query(self, census_columns, state, county, year= 2016,
                    tract=None):
        if tract is None:
            tract = "*"
        return self._query(census_columns, state, county,
                           forstr="tract:%s" % tract,
                           year=year)

    def _query(self, census_columns, state, county, forstr,
               year, tract=None):
        c = self.c

        state, county = self.try_fips_lookup(state, county)

        if tract is None:
            in_str = 'state:%s county:%s' % (state, county)
        else:
            in_str = 'state:%s county:%s tract:%s' % (state, county, tract)

        dfs = []

        # unfortunately the api only queries 50 columns at a time
        # leave room for a few extra id columns
        def chunks(l, n):
            """ Yield successive n-sized chunks from l.
            """
            for i in range(0, len(l), n):
                yield l[i:i+n]

        for census_column_batch in chunks(census_columns, 45):
            census_column_batch = list(census_column_batch)
            d = c.acs5.get(['NAME'] + census_column_batch,
                          geo={'for': forstr,
                               'in': in_str}, year=year)
            df = pd.DataFrame(d)
            df[census_column_batch] = df[census_column_batch].astype('int')
            dfs.append(df)

        assert len(dfs) >= 1
        df = dfs[0]
        for mdf in dfs[1:]:
            df = pd.merge(df, mdf, on="NAME", suffixes=("", "_ignore"))
            drop_cols = list(filter(lambda x: "_ignore" in x, df.columns))
            df = df.drop(drop_cols, axis=1)

        return df

    def block_group_and_tract_query(self, block_group_columns,
                                    tract_columns, state, county,
                                    merge_columns, block_group_size_attr,
                                    tract_size_attr, year= 2016, tract=None):
        df2 = self.tract_query(tract_columns, state, county, tract=tract,
                               year=year)
        df1 = self.block_group_query(block_group_columns, state, county,
                                     tract=tract, year=year)

        df = self._scale_and_merge(df1, block_group_size_attr, df2,
                                   tract_size_attr, tract_columns,
                                   merge_columns, suffixes=("", "_ignore"))
        drop_cols = list(filter(lambda x: "_ignore" in x, df.columns))
        df = df.drop(drop_cols, axis=1)

        return df

    def _get_pums_relationship(self):
        if self.pums_relationship_df is None:
            self.pums_relationship_df = \
                pd.read_csv(self.pums_relationship_file_url, dtype={
                    "statefp": "object",
{'fraction_non_alphanumeric': 0.0999208547685002, 'fraction_numerical': 0.013256826276216858, 'mean_word_length': 1.7925414364640884, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2944027', 'n_tokens_mistral': 3255, 'n_tokens_neox': 3100, 'n_words': 710}