<commit_before>import os
import sys
cur_path = os.path.abspath(os.path.dirname(__file__))
sys.path.append(os.path.join(cur_path, "../../"))
sys.path.append(os.path.join(cur_path, "../"))
import numpy as np
import pandas as pd
import pytest
import numpy.testing as npt
from pandas import DataFrame, Series
from pandas.util.testing import assert_frame_equal
from pandas.util.testing import assert_series_equal
from numba import jit, vectorize, guvectorize
from taxcalc import *


data = [[1.0, 2, 'a'],
        [-1.0, 4, 'a'],
        [3.0, 6, 'a'],
        [2.0, 4, 'b'],
        [3.0, 6, 'b']]

irates = {1991:0.015, 1992:0.020, 1993:0.022, 1994:0.020, 1995:0.021,
          1996:0.022, 1997:0.023, 1998:0.024, 1999:0.024, 2000:0.024,
          2001:0.024, 2002:0.024}

def test_expand_1D_short_array():
    x = np.array([4, 5, 9], dtype='i4')
    exp2 = np.array([9.0 * math.pow(1.02, i) for i in range(1, 8)])
    exp1 = np.array([4, 5, 9])
    exp = np.zeros(10)
    exp[:3] = exp1
    exp[3:] = exp2
    res = expand_1D(x, inflate=True, inflation_rates=[0.02]*10, num_years=10)
    assert(np.allclose(exp.astype(x.dtype, casting='unsafe'), res))

def test_expand_1D_variable_rates():
    x = np.array([4, 5, 9], dtype='f4')
    irates = [0.02, 0.02, 0.02, 0.03, 0.035]
    exp2 = []
    cur = 9.0
    for i in range(1, 3):
        idx = i + len(x) - 1
        cur *= (1.0 + irates[idx])
        exp2.append(cur)

    exp1 = np.array([4, 5, 9])
    exp = np.zeros(5)
    exp[:3] = exp1
    exp[3:] = exp2
    res = expand_1D(x, inflate=True, inflation_rates=irates, num_years=5)
    assert(np.allclose(exp.astype('f4', casting='unsafe'), res))


def test_expand_1D_scalar():
    x = 10.0
    exp = np.array([10.0 * math.pow(1.02, i) for i in range(0, 10)])
    res = expand_1D(x, inflate=True, inflation_rates=[0.02]*10, num_years=10)
    assert(np.allclose(exp, res))


def test_expand_1D_accept_None():
    x = [4., 5., None]
    irates = [0.02, 0.02, 0.02, 0.03, 0.035]
    exp2 = []
    cur = 5.0
    short_x = np.array([4, 5])
    for i in range(1, 4):
        idx = i + len(short_x) - 1
        cur *= (1.0 + irates[idx])
        exp2.append(cur)

    exp1 = np.array([4, 5])
    exp = np.zeros(5)
    exp[:2] = exp1
    exp[2:] = exp2
    res = expand_array(x, inflate=True, inflation_rates=irates, num_years=5)
    assert(np.allclose(exp.astype('f4', casting='unsafe'), res))


def test_expand_2D_short_array():
    x = np.array([[1, 2, 3]], dtype='f8')
    val = np.array([1, 2, 3], dtype='f8')
    exp2 = np.array([val * math.pow(1.02, i) for i in range(1, 5)])
    exp1 = np.array([1, 2, 3], dtype='f8')
    exp = np.zeros((5, 3))
    exp[:1] = exp1
    exp[1:] = exp2
    res = expand_2D(x, inflate=True, inflation_rates=[0.02]*5, num_years=5)
    assert(np.allclose(exp, res))


def test_expand_2D_variable_rates():
    x = np.array([[1, 2, 3]], dtype='f8')
    cur = np.array([1, 2, 3], dtype='f8')
    irates = [0.02, 0.02, 0.02, 0.03, 0.035]

    exp2 = []
    for i in range(0, 4):
        idx = i + len(x) - 1
        cur = np.array(cur*(1.0 + irates[idx]))
        print("cur is ", cur)
        exp2.append(cur)

    #exp2 = np.array([val * math.pow(1.02, i) for i in range(1, 5)])
    exp1 = np.array([1, 2, 3], dtype='f8')
    exp = np.zeros((5, 3))
    exp[:1] = exp1
    exp[1:] = exp2
    res = expand_2D(x, inflate=True, inflation_rates=irates, num_years=5)
    npt.assert_array_equal(res, np.array(exp).astype('f8', casting='unsafe'))
    assert(np.allclose(exp, res))



def test_create_tables():
    # Default Plans
    #Create a Public Use File object
    cur_path = os.path.abspath(os.path.dirname(__file__))
    tax_dta_path = os.path.join(cur_path, "../../tax_all1991_puf.gz")
    # Create a default Parameters object
    params1 = Parameters(start_year=1991, inflation_rates=irates)
    records1 = Records(tax_dta_path)
    # Create a Calculator
    calc1 = Calculator(parameters=params1, records=records1)
    calc1.calc_all()

    # User specified Plans
    user_mods = '{"1991": {"_II_rt4": [0.56]}}'
    params2 = Parameters(start_year=1991, inflation_rates=irates)
    records2 = Records(tax_dta_path)
    # Create a Calculator
    calc2 = calculator(parameters=params2, records=records2, mods=user_mods)

    calc2.calc_all()

    t2 = create_distribution_table(calc2, groupby="small_agi_bins", result_type = "weighted_sum")
    #make large agi bins table
    tdiff = create_difference_table(calc1, calc2, groupby="large_agi_bins")
    #make webapp agi bins table
    tdiff_webapp = create_difference_table(calc1, calc2, groupby="webapp_agi_bins")


def test_weighted_count_lt_zero():
    df = DataFrame(data=data, columns=['tax_diff', 's006', 'label'])
    grped = df.groupby('label')
    diffs = grped.apply(weighted_count_lt_zero, 'tax_diff')
    exp = Series(data=[4, 0], index=['a', 'b'])
    assert_series_equal(exp, diffs)


def test_weighted_count_gt_zero():
    df = DataFrame(data=data, columns=['tax_diff', 's006', 'label'])
    grped = df.groupby('label')
    diffs = grped.apply(weighted_count_gt_zero, 'tax_diff')
    exp = Series(data=[8, 10], index=['a', 'b'])
    assert_series_equal(exp, diffs)
 

def test_weighted_count():
    df = DataFrame(data=data, columns=['tax_diff', 's006', 'label'])
    grped = df.groupby('label')
    diffs = grped.apply(weighted_count)
    exp = Series(data=[12, 10], index=['a', 'b'])
    assert_series_equal(exp, diffs)
 

def test_weighted_mean():
    df = DataFrame(data=data, columns=['tax_diff', 's006', 'label'])
    grped = df.groupby('label')
    diffs = grped.apply(weighted_mean, 'tax_diff')
    exp = Series(data=[16.0/12.0, 26.0/10.0], index=['a', 'b'])
    assert_series_equal(exp, diffs)
 

def test_weighted_sum():
    df = DataFrame(data=data, columns=['tax_diff', 's006', 'label'])
    grped = df.groupby('label')
    diffs = grped.apply(weighted_sum, 'tax_diff')
    exp = Series(data=[16.0, 26.0], index=['a', 'b'])
    assert_series_equal(exp, diffs)
 
<commit_msg>Added test for float comparison in weighted_count_gt/lt in test_utils.py<commit_after>import os
import sys
cur_path = os.path.abspath(os.path.dirname(__file__))
sys.path.append(os.path.join(cur_path, "../../"))
sys.path.append(os.path.join(cur_path, "../"))
import numpy as np
import pandas as pd
import pytest
import numpy.testing as npt
from pandas import DataFrame, Series
from pandas.util.testing import assert_frame_equal
from pandas.util.testing import assert_series_equal
from numba import jit, vectorize, guvectorize
from taxcalc import *


data = [[1.0, 2, 'a'],
        [-1.0, 4, 'a'],
        [3.0, 6, 'a'],
        [2.0, 4, 'b'],
        [3.0, 6, 'b']]

data_float = [[1.0, 2, 'a'],
              [-1.0, 4, 'a'],
              [0.0000000001, 3, 'a'],
              [-0.0000000001, 1, 'a'],
              [3.0, 6, 'a'],
              [2.0, 4, 'b'],
              [0.0000000001, 3, 'b'],
              [-0.0000000001, 1, 'b'],
              [3.0, 6, 'b']]

irates = {1991:0.015, 1992:0.020, 1993:0.022, 1994:0.020, 1995:0.021,
          1996:0.022, 1997:0.023, 1998:0.024, 1999:0.024, 2000:0.024,
          2001:0.024, 2002:0.024}

def test_expand_1D_short_array():
    x = np.array([4, 5, 9], dtype='i4')
    exp2 = np.array([9.0 * math.pow(1.02, i) for i in range(1, 8)])
    exp1 = np.array([4, 5, 9])
    exp = np.zeros(10)
    exp[:3] = exp1
    exp[3:] = exp2
    res = expand_1D(x, inflate=True, inflation_rates=[0.02]*10, num_years=10)
    assert(np.allclose(exp.astype(x.dtype, casting='unsafe'), res))

def test_expand_1D_variable_rates():
    x = np.array([4, 5, 9], dtype='f4')
    irates = [0.02, 0.02, 0.02, 0.03, 0.035]
    exp2 = []
    cur = 9.0
    for i in range(1, 3):
        idx = i + len(x) - 1
        cur *= (1.0 + irates[idx])
        exp2.append(cur)

    exp1 = np.array([4, 5, 9])
    exp = np.zeros(5)
    exp[:3] = exp1
    exp[3:] = exp2
    res = expand_1D(x, inflate=True, inflation_rates=irates, num_years=5)
    assert(np.allclose(exp.astype('f4', casting='unsafe'), res))


def test_expand_1D_scalar():
    x = 10.0
    exp = np.array([10.0 * math.pow(1.02, i) for i in range(0, 10)])
    res = expand_1D(x, inflate=True, inflation_rates=[0.02]*10, num_years=10)
    assert(np.allclose(exp, res))


def test_expand_1D_accept_None():
    x = [4., 5., None]
    irates = [0.02, 0.02, 0.02, 0.03, 0.035]
    exp2 = []
    cur = 5.0
    short_x = np.array([4, 5])
    for i in range(1, 4):
        idx = i + len(short_x) - 1
        cur *= (1.0 + irates[idx])
        exp2.append(cur)

    exp1 = np.array([4, 5])
    exp = np.zeros(5)
    exp[:2] = exp1
    exp[2:] = exp2
    res = expand_array(x, inflate=True, inflation_rates=irates, num_years=5)
    assert(np.allclose(exp.astype('f4', casting='unsafe'), res))


def test_expand_2D_short_array():
    x = np.array([[1, 2, 3]], dtype='f8')
    val = np.array([1, 2, 3], dtype='f8')
    exp2 = np.array([val * math.pow(1.02, i) for i in range(1, 5)])
    exp1 = np.array([1, 2, 3], dtype='f8')
    exp = np.zeros((5, 3))
    exp[:1] = exp1
    exp[1:] = exp2
    res = expand_2D(x, inflate=True, inflation_rates=[0.02]*5, num_years=5)
    assert(np.allclose(exp, res))


def test_expand_2D_variable_rates():
    x = np.array([[1, 2, 3]], dtype='f8')
    cur = np.array([1, 2, 3], dtype='f8')
    irates = [0.02, 0.02, 0.02, 0.03, 0.035]

    exp2 = []
    for i in range(0, 4):
        idx = i + len(x) - 1
        cur = np.array(cur*(1.0 + irates[idx]))
        print("cur is ", cur)
        exp2.append(cur)

    #exp2 = np.array([val * math.pow(1.02, i) for i in range(1, 5)])
    exp1 = np.array([1, 2, 3], dtype='f8')
    exp = np.zeros((5, 3))
    exp[:1] = exp1
    exp[1:] = exp2
    res = expand_2D(x, inflate=True, inflation_rates=irates, num_years=5)
    npt.assert_array_equal(res, np.array(exp).astype('f8', casting='unsafe'))
    assert(np.allclose(exp, res))



def test_create_tables():
    # Default Plans
    #Create a Public Use File object
    cur_path = os.path.abspath(os.path.dirname(__file__))
    tax_dta_path = os.path.join(cur_path, "../../tax_all1991_puf.gz")
    # Create a default Parameters object
    params1 = Parameters(start_year=1991, inflation_rates=irates)
    records1 = Records(tax_dta_path)
    # Create a Calculator
    calc1 = Calculator(parameters=params1, records=records1)
    calc1.calc_all()

    # User specified Plans
    user_mods = '{"1991": {"_II_rt4": [0.56]}}'
    params2 = Parameters(start_year=1991, inflation_rates=irates)
    records2 = Records(tax_dta_path)
    # Create a Calculator
    calc2 = calculator(parameters=params2, records=records2, mods=user_mods)

    calc2.calc_all()

    t2 = create_distribution_table(calc2, groupby="small_agi_bins", result_type = "weighted_sum")
    #make large agi bins table
    tdiff = create_difference_table(calc1, calc2, groupby="large_agi_bins")
    #make webapp agi bins table
    tdiff_webapp = create_difference_table(calc1, calc2, groupby="webapp_agi_bins")


def test_weighted_count_lt_zero():
    df1 = DataFrame(data=data, columns=['tax_diff', 's006', 'label'])
    grped = df1.groupby('label')
    diffs = grped.apply(weighted_count_lt_zero, 'tax_diff')
    exp = Series(data=[4, 0], index=['a', 'b'])
    assert_series_equal(exp, diffs)
    df2 = DataFrame(data=data_float, columns=['tax_diff', 's006', 'label'])
    grped = df2.groupby('label')
    diffs = grped.apply(weighted_count_lt_zero, 'tax_diff')
    exp = Series(data=[4, 0], index=['a', 'b'])
    assert_series_equal(exp, diffs)


def test_weighted_count_gt_zero():
    df1 = DataFrame(data=data, columns=['tax_diff', 's006', 'label'])
    grped = df1.groupby('label')
    diffs = grped.apply(weighted_count_gt_zero, 'tax_diff')
    exp = Series(data=[8, 10], index=['a', 'b'])
    assert_series_equal(exp, diffs)
    df2 = DataFrame(data=data, columns=['tax_diff', 's006', 'label'])
    grped = df2.groupby('label')
    diffs = grped.apply(weighted_count_gt_zero, 'tax_diff')
    exp = Series(data=[8, 10], index=['a', 'b'])
    assert_series_equal(exp, diffs)
 

def test_weighted_count():
    df = DataFrame(data=data, columns=['tax_diff', 's006', 'label'])
    grped = df.groupby('label')
    diffs = grped.apply(weighted_count)
    exp = Series(data=[12, 10], index=['a', 'b'])
    assert_series_equal(exp, diffs)
 

def test_weighted_mean():
    df = DataFrame(data=data, columns=['tax_diff', 's006', 'label'])
    grped = df.groupby('label')
    diffs = grped.apply(weighted_mean, 'tax_diff')
    exp = Series(data=[16.0/12.0, 26.0/10.0], index=['a', 'b'])
    assert_series_equal(exp, diffs)
 

def test_weighted_sum():
    df = DataFrame(data=data, columns=['tax_diff', 's006', 'label'])
    grped = df.groupby('label')
    diffs = grped.apply(weighted_sum, 'tax_diff')
    exp = Series(data=[16.0, 26.0], index=['a', 'b'])
    assert_series_equal(exp, diffs)
 

{'fraction_non_alphanumeric': 0.17843489116017788, 'fraction_numerical': 0.07021924007177967, 'mean_word_length': 3.3377326565143823, 'pattern_counts': {'":': 4, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 4, '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': '93771', 'n_tokens_mistral': 5629, 'n_tokens_neox': 5040, 'n_words': 1064}