<commit_before>#!/usr/bin/env python
"""
Pull together the NCES data and the Segragation index calculator to produce
some reports on segregation in the USA.
"""
import sys
import argparse

from segcalc import SegCalc
from nces_parser import NCESParser

from xlwt import Workbook

# ==============================================================================
# Constants
# ==============================================================================
fips_to_st = {
    1: ('AL', 'Alabama'),
    2: ('AK', 'Alaska'),
    4: ('AZ', 'Arizona'),
    5: ('AR', 'Arkansas'),
    6: ('CA', 'California'),
    8: ('CO', 'Colorado'),
    9: ('CT', 'Connecticut'),
    10: ('DE', 'Delaware'),
    11: ('DC', 'District of Columbia'),
    12: ('FL', 'Florida'),
    13: ('GA', 'Georgia'),
    15: ('HI', 'Hawaii'),
    16: ('ID', 'Idaho'),
    17: ('IL', 'Illinois'),
    18: ('IN', 'Indiana'),
    19: ('IA', 'Iowa'),
    20: ('KS', 'Kansas'),
    21: ('KY', 'Kentucky'),
    22: ('LA', 'Louisiana'),
    23: ('ME', 'Maine'),
    24: ('MD', 'Maryland'),
    25: ('MA', 'Massachusetts'),
    26: ('MI', 'Michigan'),
    27: ('MN', 'Minnesota'),
    28: ('MS', 'Mississippi'),
    29: ('MO', 'Missouri'),
    30: ('MT', 'Montana'),
    31: ('NE', 'Nebraska'),
    32: ('NV', 'Nevada'),
    33: ('NH', 'New Hampshire'),
    34: ('NJ', 'New Jersey'),
    35: ('NM', 'New Mexico'),
    36: ('NY', 'New York'),
    37: ('NC', 'North Carolina'),
    38: ('ND', 'North Dakota'),
    39: ('OH', 'Ohio'),
    40: ('OK', 'Oklahoma'),
    41: ('OR', 'Oregon'),
    42: ('PA', 'Pennsylvania'),
    44: ('RI', 'Rhode Island'),
    45: ('SC', 'South Carolina'),
    46: ('SD', 'South Dakota'),
    47: ('TN', 'Tennessee'),
    48: ('TX', 'Texas'),
    49: ('UT', 'Utah'),
    50: ('VT', 'Vermont'),
    51: ('VA', 'Virginia'),
    53: ('WA', 'Washington'),
    54: ('WV', 'West Virginia'),
    55: ('WI', 'Wisconsin'),
    56: ('WY', 'Wyoming')
}


# ==============================================================================
# Functions
# ==============================================================================
def calc_idxes(year_range, idx):
    """
    Call down to get all the various measures calculated
    """
    exp_idx = []
    iso_idx = []
    dis_idx = []
    for year in year_range:
        print "Working on NCES Data from:  %d" % year
        # Get our dataset ready
        nces = NCESParser(year=year)
        schools = nces.parse(make_dict=True)
        sg = SegCalc(schools, idx)
        exp_idx.append(sg.calc_exp_idx())
        iso_idx.append(sg.calc_iso_idx())
        dis_idx.append(sg.calc_dis_idx())
    return (exp_idx, iso_idx, dis_idx)

# -------------------------------------
def save_report(year_range, idxes, filename):
    # Break out the calculated data
    exp_idx, iso_idx, dis_idx = idxes

    # Now lets format it into a report.
    # Report will be a 2D matrix
    # X-Axis = school year
    # Y-Axis = FIPS Code
    wb = Workbook()
    ews = wb.add_sheet('Exposure Index')
    iws = wb.add_sheet('Isolation Index')
    dws = wb.add_sheet('Dissimilarity Index')

    worksheets = [ews, iws, dws]

    # Create the headers/labels row/col
    for ws in worksheets:
        ws.write(0, 0, "State/Year")
        for j, st in enumerate(fips_to_st):
            ws.write(j+1, 0, fips_to_st[st][0])

    # Print out the data
    for i, year in enumerate(year_range):
        print "Write Report for:  %d" % year
        for ws in worksheets:
            ws.write(0, i+1, year)
        for j, st in enumerate(fips_to_st.keys()):
            if exp_idx[i][st] < 0.001:
                ews.write(j+1, i+1, "")
            else:
                ews.write(j+1, i+1, exp_idx[i][st])

            if iso_idx[i][st] < 0.001:
                iws.write(j+1, i+1, "")
            else:
                iws.write(j+1, i+1, iso_idx[i][st])

            if dis_idx[i][st] < 0.001:
                dws.write(j+1, i+1, "")
            else:
                dws.write(j+1, i+1, dis_idx[i][st])
    wb.save(filename)

# -------------------------------------
# Parse the command line options
# -------------------------------------
def main(argv):
    parser = argparse.ArgumentParser(description='Segregation Report Generator')
    parser.add_argument('--year', action='store', dest='year', required=False, type=int,
            help='NCES Data Year')
    parser.add_argument('--outfile', action='store', dest='outfile', required=False,
            help='Report File')
    args = parser.parse_args()

    # Lets calculate all the data first
    year_range = range(1987, 2011)

    """
    idx = {'Y_GROUP': 'BLACK', 'Z_GROUP': 'WHITE', 'TOTAL': 'MEMBER', 'CATEGORY': 'FIPS', 'SUB_CAT': 'LEAID'}
    idxes = calc_idxes(year_range, idx)
    save_report(year_range, idxes, "black_" + args.outfile)

    idx = {'Y_GROUP': 'HISP', 'Z_GROUP': 'WHITE', 'TOTAL': 'MEMBER', 'CATEGORY': 'FIPS', 'SUB_CAT': 'LEAID'}
    idxes = calc_idxes(year_range, idx)
    save_report(year_range, idxes, "hisp_" + args.outfile)

    idx = {'Y_GROUP': 'ASIAN', 'Z_GROUP': 'WHITE', 'TOTAL': 'MEMBER', 'CATEGORY': 'FIPS', 'SUB_CAT': 'LEAID'}
    idxes = calc_idxes(year_range, idx)
    save_report(year_range, idxes, "asian_" + args.outfile)
    """

    idx = {'Y_GROUP': 'AM', 'Z_GROUP': 'WHITE', 'TOTAL': 'MEMBER', 'CATEGORY': 'FIPS', 'SUB_CAT': 'LEAID'}
    idxes = calc_idxes(year_range, idx)
    save_report(year_range, idxes, "native_amer_" + args.outfile)

# -------------------------------------
# Drop the script name from the args
# and call our command line parser
# -------------------------------------
if __name__ == "__main__":
    main(sys.argv[1:])



<commit_msg>Modified to allow a broader set of analysis options:<commit_after>#!/usr/bin/env python
"""
Pull together the NCES data and the Segragation index calculator to produce
some reports on segregation in the USA.
"""
import sys
import argparse

from segcalc import SegCalc
from nces_parser import NCESParser

from xlwt import Workbook

# ==============================================================================
# Constants
# ==============================================================================
fips_to_st = {
    1: ('AL', 'Alabama'),
    2: ('AK', 'Alaska'),
    4: ('AZ', 'Arizona'),
    5: ('AR', 'Arkansas'),
    6: ('CA', 'California'),
    8: ('CO', 'Colorado'),
    9: ('CT', 'Connecticut'),
    10: ('DE', 'Delaware'),
    11: ('DC', 'District of Columbia'),
    12: ('FL', 'Florida'),
    13: ('GA', 'Georgia'),
    15: ('HI', 'Hawaii'),
    16: ('ID', 'Idaho'),
    17: ('IL', 'Illinois'),
    18: ('IN', 'Indiana'),
    19: ('IA', 'Iowa'),
    20: ('KS', 'Kansas'),
    21: ('KY', 'Kentucky'),
    22: ('LA', 'Louisiana'),
    23: ('ME', 'Maine'),
    24: ('MD', 'Maryland'),
    25: ('MA', 'Massachusetts'),
    26: ('MI', 'Michigan'),
    27: ('MN', 'Minnesota'),
    28: ('MS', 'Mississippi'),
    29: ('MO', 'Missouri'),
    30: ('MT', 'Montana'),
    31: ('NE', 'Nebraska'),
    32: ('NV', 'Nevada'),
    33: ('NH', 'New Hampshire'),
    34: ('NJ', 'New Jersey'),
    35: ('NM', 'New Mexico'),
    36: ('NY', 'New York'),
    37: ('NC', 'North Carolina'),
    38: ('ND', 'North Dakota'),
    39: ('OH', 'Ohio'),
    40: ('OK', 'Oklahoma'),
    41: ('OR', 'Oregon'),
    42: ('PA', 'Pennsylvania'),
    44: ('RI', 'Rhode Island'),
    45: ('SC', 'South Carolina'),
    46: ('SD', 'South Dakota'),
    47: ('TN', 'Tennessee'),
    48: ('TX', 'Texas'),
    49: ('UT', 'Utah'),
    50: ('VT', 'Vermont'),
    51: ('VA', 'Virginia'),
    53: ('WA', 'Washington'),
    54: ('WV', 'West Virginia'),
    55: ('WI', 'Wisconsin'),
    56: ('WY', 'Wyoming')
}


# Maximum number of rows to report in an output file.
MAX_REPORT = 200

# ==============================================================================
# Functions
# ==============================================================================
def calc_idxes(segcalc):
    """
    Call down to get all the various measures calculated
    """
    exp_idx = segcalc.calc_exp_idx()
    iso_idx = segcalc.calc_iso_idx()
    dis_idx = segcalc.calc_dis_idx()
    tot_idx = segcalc.calc_totals()
    return (exp_idx, iso_idx, dis_idx, tot_idx)

# --------------------------------------
def calc_idxes_range(year_range, idx):
    """
    Get all the results over the years
    """
    exp_idx = []
    iso_idx = []
    dis_idx = []
    for year in year_range:
        print "Working on NCES Data from:  %d" % year
        # Get our dataset ready
        nces = NCESParser(year=year)
        schools = nces.parse(make_dict=True)
        sg = SegCalc(schools, idx)
        exp_idx.append(sg.calc_exp_idx())
        iso_idx.append(sg.calc_iso_idx())
        dis_idx.append(sg.calc_dis_idx())
    return (exp_idx, iso_idx, dis_idx)

# -------------------------------------
def save_report(year_range, idxes, category_list, category_txt, filename):
    """
    Write out a bunch of report data to a spreadsheet report.
    Report will be a 2D matrix:
        - X-Axis = school year
        - Y-Axis = 'Category' (FIPS Code, District, etc...)

    Notes:
        - idxes contains the data
        - worksheets is a list of XLS worksheets, one per report in idxes
    """
    wb = Workbook()
    ews = wb.add_sheet('Exposure Index')
    iws = wb.add_sheet('Isolation Index')
    dws = wb.add_sheet('Dissimilarity Index')
    size = wb.add_sheet('Student Count')

    worksheets = [ews, iws, dws, size]

    # Create the headers/labels row/col
    for ws in worksheets:
        ws.write(0, 0, "LEA/Year")
        for j, st in enumerate(category_list):
            if j < MAX_REPORT:
                if len(category_txt[st]) == 2: # Don't change caps for State abbr.
                    ws.write(j+1, 0, category_txt[st])
                else:
                    ws.write(j+1, 0, category_txt[st].title())

    # Print out the data
    for i, year in enumerate(year_range):
        print "Write Report for:  %d" % year
        for ws in worksheets:
            ws.write(0, i+1, year)
        for j, st in enumerate(category_list):
            if j < MAX_REPORT:
                for k, idx in enumerate(idxes):
                    try:
                        if idx[i][st] < 0.001:
                            worksheets[k].write(j+1, i+1, "")
                        else:
                            worksheets[k].write(j+1, i+1, idx[i][st])
                    except KeyError:
                        worksheets[k].write(j+1, i+1, "")
    wb.save(filename)

# -------------------------------------
# Parse the command line options
# -------------------------------------
def main(argv):
    parser = argparse.ArgumentParser(description='Segregation Report Generator')
    parser.add_argument('--year', action='store', dest='year', required=False, type=int,
            help='NCES Data Year')
    parser.add_argument('--outfile', action='store', dest='outfile', required=True,
            help='Report File')
    parser.add_argument('--category', action='store', dest='category', required=False,
            help='Which Category do we sort the results by?')
    parser.add_argument('--match_idx', action='store', dest='match_idx', required=False,
            help='Only use data points that match some criterion')
    parser.add_argument('--match_val', action='store', dest='match_val', required=False,
            help='Only use data points that match some criterion')
    parser.add_argument('-debug', action='store_true', dest='debug', required=False,
            help='Debug Mode')
    args = parser.parse_args()

    if args.category:
        category = args.category
    else:
        category = 'FIPS'

    # Lets calculate all the data first
    if args.debug:
        year_range = [1987]
        groups = ['BLACK']
    else:
        year_range = range(1987, 2011)
        groups = ['BLACK', 'HISP', 'ASIAN', 'AM']

    # Default search query
    idx = {
        'Y_GROUP': "BLACK",
        'Z_GROUP': 'WHITE',
        'TOTAL': 'MEMBER',
        'CATEGORY': category,
        'SUB_CAT': 'LEAID',
    }
    if args.match_idx:
        idx['MATCH_IDX'] = args.match_idx
        idx['MATCH_VAL'] = args.match_val

    idxes = [[], [], [], []]
    for group in groups:
        for year in year_range:
            print "Loading NCES Data from:  %d" % year
            nces = NCESParser(year=year)
            schools = nces.parse(make_dict=True)
            # Get our data query ready
            idx['Y_GROUP'] = group
            segcalc = SegCalc(schools, idx)
            if category == 'LEAID':
                category_lut = segcalc.get_idxed_val('LEAID', 'LEANM')
            elif category == 'FIPS':
                category_lut = dict(zip(fips_to_st.keys(), [fips_to_st[key][0] for key in fips_to_st.keys()]))

            exp,iso,dis,tot = calc_idxes(segcalc)
            category_by_size = sorted(tot, key=tot.get, reverse=True)
            # Filter out keys absent from our report tables.
            category_by_size = [i for i in category_by_size if i in category_lut.keys()]

            idxes[0].append(exp)
            idxes[1].append(iso)
            idxes[2].append(dis)
            idxes[3].append(tot)

        save_report(year_range, idxes, category_by_size, category_lut, group.lower() + '_' + args.outfile)

# -------------------------------------
# Drop the script name from the args
# and call our command line parser
# -------------------------------------
if __name__ == "__main__":
    main(sys.argv[1:])




{'fraction_non_alphanumeric': 0.2351281099145934, 'fraction_numerical': 0.02020051986632009, 'mean_word_length': 2.4528205128205127, 'pattern_counts': {'":': 2, '<': 9, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5029074', 'n_tokens_mistral': 4663, 'n_tokens_neox': 4304, 'n_words': 1155}