def test_get_e_orig_exception2():
'''
Test that RuntimeError is raised if abil_wgts does not have length 87
'''
age_wgts = np.ones(80) * 1 / 80
abil_wgts = np.array([0.25, 0.25, 0.2, 0.1, 0.1])
with pytest.raises(RuntimeError):
income.get_e_orig(age_wgts, abil_wgts)
def test_plot_income_data():
ages = np.linspace(20 + 0.5, 100 - 0.5, 80)
abil_midp = np.array([0.125, 0.375, 0.6, 0.75, 0.85, 0.945, 0.995])
abil_pcts = np.array([0.25, 0.25, 0.2, 0.1, 0.1, 0.09, 0.01])
age_wgts = np.ones(80) * 1 / 80
emat = income.get_e_orig(age_wgts, abil_pcts)
fig = parameter_plots.plot_income_data(
ages, abil_midp, abil_pcts, emat)
assert fig
def test_plot_income_data_save_fig(tmpdir):
ages = np.linspace(20 + 0.5, 100 - 0.5, 80)
abil_midp = np.array([0.125, 0.375, 0.6, 0.75, 0.85, 0.945, 0.995])
abil_pcts = np.array([0.25, 0.25, 0.2, 0.1, 0.1, 0.09, 0.01])
age_wgts = np.ones(80) * 1 / 80
emat = income.get_e_orig(age_wgts, abil_pcts)
parameter_plots.plot_income_data(
ages, abil_midp, abil_pcts, emat, output_dir=tmpdir)
img1 = mpimg.imread(os.path.join(tmpdir, 'ability_3D_lev.png'))
img2 = mpimg.imread(os.path.join(tmpdir, 'ability_3D_log.png'))
img3 = mpimg.imread(os.path.join(tmpdir, 'ability_2D_log.png'))
assert isinstance(img1, np.ndarray)
assert isinstance(img2, np.ndarray)
assert isinstance(img3, np.ndarray)
def test_get_e_orig():
'''
Test get_e_orig() function
'''
expected_vals = np.array([
[0.41637813, 0.27892939, 0.29770352, 0.34999537, 0.43550237,
0.7104421, 1.84111771],
[0.40966052, 0.29351488, 0.32491809, 0.38673085, 0.48142028,
0.75826254, 1.99770813],
[0.40406881, 0.30862551, 0.35340211, 0.42551707, 0.5298768,
0.80925704, 2.16483317],
[0.39951541, 0.32425432, 0.38308642, 0.46625585, 0.580742,
0.86355704, 2.3427626],
[0.39592303, 0.34039127, 0.41388897, 0.50882471, 0.63385587,
0.92128697, 2.53170513],
[0.39322307, 0.35702297, 0.44571523, 0.55307772, 0.68902964,
0.98256181, 2.73179742],
[0.39135441, 0.37413249, 0.47845898, 0.59884672, 0.7460476,
1.04748443, 2.94309262],
[0.3902622, 0.39169918, 0.51200315, 0.645943, 0.80466953,
1.11614268, 3.16554852],
[0.38989688, 0.40969841, 0.54622089, 0.69415939, 0.8646336,
1.18860621, 3.39901567],
[0.39021328, 0.42810145, 0.58097684, 0.74327267, 0.92565968,
1.26492305, 3.64322541],
[0.39116983, 0.44687529, 0.61612841, 0.79304632, 0.98745294,
1.34511602, 3.89777836],
[0.39272783, 0.46598253, 0.65152732, 0.84323347, 1.04970783,
1.42917894, 4.16213344],
[0.3948508, 0.48538127, 0.68702111, 0.89358009, 1.11211214,
1.51707268, 4.43559781],
[0.3975039, 0.50502501, 0.72245476, 0.94382825, 1.17435118,
1.60872119, 4.71731806],
[0.40065336, 0.52486267, 0.75767235, 0.9937194, 1.23611194,
1.70400743, 5.00627307],
[0.404266, 0.54483855, 0.79251866, 1.04299768, 1.29708717,
1.80276946, 5.30126881],
[0.40830873, 0.56489242, 0.8268408, 1.09141311, 1.35697929,
1.90479656, 5.60093551],
[0.41274816, 0.58495958, 0.86048981, 1.13872463, 1.41550405,
2.00982572, 5.90372761],
[0.41755017, 0.60497105, 0.89332205, 1.1847029, 1.47239387,
2.11753837, 6.20792673],
[0.42267953, 0.62485373, 0.9252007, 1.22913293, 1.5274008,
2.22755771, 6.51164811],
[0.42809957, 0.6445307, 0.95599693, 1.27181632, 1.58029914,
2.33944652, 6.81285071],
[0.4337719, 0.66392152, 0.98559107, 1.31257321, 1.63088754,
2.45270577, 7.10935121],
[0.43965608, 0.68294259, 1.0138736, 1.35124392, 1.6789907,
2.56677413, 7.39884199],
[0.44570939, 0.70150761, 1.04074596, 1.38769019, 1.72446061,
2.68102845, 7.67891322],
[0.45188667, 0.71952807, 1.06612125, 1.42179604, 1.76717729,
2.79478531, 7.947079],
[0.45814012, 0.73691379, 1.08992465, 1.45346833, 1.80704915,
2.90730394, 8.20080721],
[0.46441919, 0.75357353, 1.11209386, 1.4826369, 1.8440129,
3.01779041, 8.43755302],
[0.47067057, 0.76941566, 1.13257918, 1.50925445, 1.87803307,
3.12540323, 8.65479541],
[0.47683814, 0.78434886, 1.15134357, 1.5332961, 1.90910121,
3.22926052, 8.85007629],
[0.48286305, 0.79828285, 1.16836252, 1.5547586, 1.93723475,
3.32844853, 9.02104147],
[0.48868388, 0.81112921, 1.18362376, 1.57365942, 1.96247564,
3.42203181, 9.16548263],
[0.49423684, 0.82280213, 1.19712692, 1.59003557, 1.98488875,
3.50906461, 9.28137951],
[0.49945604, 0.83321932, 1.20888294, 1.60394223, 2.00456011,
3.58860379, 9.36694123],
[0.50427389, 0.84230278, 1.21891355, 1.61545125, 2.02159501,
3.65972266, 9.42064568],
[0.50862155, 0.84997968, 1.22725057, 1.62464965, 2.0361161,
3.72152596, 9.44127599],
[0.51242946, 0.85618315, 1.23393515, 1.63163789, 2.04826135,
3.77316537, 9.4279529],
[0.515628, 0.86085312, 1.23901697, 1.63652827, 2.0581821,
3.81385546, 9.38016206],
[0.51814816, 0.86393709, 1.24255344, 1.63944323, 2.06604115,
3.8428897, 9.29777524],
[0.51992236, 0.8653908, 1.24460882, 1.64051371, 2.07201085,
3.85965604, 9.18106471],
[0.5208853, 0.86517897, 1.24525336, 1.63987758, 2.07627136,
3.86365177, 9.03070999],
[0.52097488, 0.86327583, 1.24456245, 1.63767811, 2.07900905,
3.8544972, 8.84779645],
[0.52013315, 0.85966566, 1.24261577, 1.63406257, 2.08041495,
3.83194769, 8.63380567],
[0.51830734, 0.85434321, 1.23949649, 1.62918093, 2.08068349,
3.79590369, 8.39059721],
[0.51545085, 0.847314, 1.23529046, 1.62318461, 2.08001127,
3.74641832, 8.12038226],
[0.51152426, 0.83859455, 1.23008548, 1.6162255, 2.07859611,
3.68370224, 7.82568944],
[0.50649633, 0.82821245, 1.22397061, 1.60845493, 2.07663628,
3.60812544, 7.50932356],
[0.50034486, 0.81620635, 1.21703552, 1.60002298, 2.07432981,
3.52021585, 7.17431825],
[0.49305761, 0.80262576, 1.20936991, 1.59107771, 2.07187416,
3.42065451, 6.82388358],
[0.484633, 0.78753081, 1.20106298, 1.58176471, 2.06946591,
3.31026739, 6.46135004],
[0.47508069, 0.77099184, 1.19220294, 1.57222669, 2.06730075,
3.19001395, 6.09011028],
[0.46442209, 0.75308881, 1.18287666, 1.5626032, 2.0655736,
3.06097248, 5.71356024],
[0.45269057, 0.73391069, 1.17316925, 1.55303051, 2.06447895,
2.92432273, 5.33504111],
[0.43993156, 0.7135547, 1.16316384, 1.54364158, 2.06421135,
2.78132605, 4.95778367],
[0.4262024, 0.6921254, 1.1529413, 1.53456619, 2.0649661,
2.63330357, 4.5848566],
[0.41157194, 0.66973375, 1.14258011, 1.52593111, 2.06694014,
2.48161309, 4.21911985],
[0.39611995, 0.64649606, 1.13215622, 1.51786046, 2.07033309,
2.32762508, 3.86318429],
[0.37993629, 0.62253291, 1.12174297, 1.51047609, 2.07534856,
2.17269864, 3.5193786],
[0.36311981, 0.59796794, 1.11141112, 1.50389813, 2.08219555,
2.01815799, 3.18972386],
[0.34577708, 0.57292671, 1.10122884, 1.49824561, 2.09109021,
1.8652701, 2.87591642],
[0.32802095, 0.54753544, 1.09126177, 1.49363716, 2.10225772,
1.71522418, 2.57931895],
[0.31097577, 0.52299945, 1.08097345, 1.48938039, 2.05910307,
1.68755637, 2.4705595],
[0.29553283, 0.50031601, 1.0697088, 1.48465828, 2.01373057,
1.65113303, 2.35125468],
[0.2814869, 0.47931181, 1.05734135, 1.4793911, 1.9661704,
1.61624875, 2.24294162],
[0.26866459, 0.45983004, 1.04372664, 1.47348009, 1.91649605,
1.58280799, 2.14416819],
[0.25691877, 0.44172934, 1.0287006, 1.46680151, 1.86483009,
1.55072298, 2.05372724],
[0.24612404, 0.42488265, 1.0120785, 1.45919837, 1.8113483,
1.51991291, 1.97060708],
[0.23617305, 0.40917598, 0.99365525, 1.45046877, 1.75628115,
1.49030327, 1.8939534],
[0.22697348, 0.3945071, 0.97320797, 1.44034904, 1.69991187,
1.46182523, 1.82303985],
[0.21844559, 0.38078435, 0.95050234, 1.42848901, 1.64257072,
1.43441515, 1.75724494],
[0.21052026, 0.36792557, 0.92530475, 1.414415, 1.58462523,
1.40801405, 1.69603377],
[0.20313726, 0.35585701, 0.89740247, 1.397473, 1.5264671,
1.38256724, 1.63894345],
[0.19624397, 0.34451244, 0.86663412, 1.37674133, 1.46849672,
1.35802389, 1.58557141],
[0.18979419, 0.33383231, 0.83293101, 1.35088923, 1.41110679,
1.33433672, 1.53556584],
[0.18374728, 0.323763, 0.79636643, 1.31795807, 1.35466677,
1.31146171, 1.48861798],
[0.17806727, 0.31425616, 0.75720385, 1.27502229, 1.29950956,
1.28935779, 1.44445568],
[0.17272232, 0.30526813, 0.71592857, 1.21773589, 1.24592155,
1.26798661, 1.40283818],
[0.16768408, 0.29675942, 0.67324352, 1.1400079, 1.19413659,
1.24731233, 1.36355167],
[0.16292728, 0.28869426, 0.63001667, 1.03495944, 1.14433385,
1.22730141, 1.32640565],
[0.1584293, 0.28104019, 0.58718372, 0.90005352, 1.09663915,
1.20792244, 1.29122983],
[0.15416985, 0.27376772, 0.54563088, 0.74681858, 1.05112886,
1.18914593, 1.25787151]])
age_wgts = np.ones(80) * 1 / 80
abil_wgts = np.array([0.25, 0.25, 0.2, 0.1, 0.1, 0.09, 0.014])
test_vals = income.get_e_orig(age_wgts, abil_wgts, plot=True)
assert np.allclose(test_vals, expected_vals)
