test_graphics.py 文件源码

def test_plot_heatmap1():
    filename = abspath(join(testdir, 'plot_heatmap1.png'))
    if isfile(filename):
        os.remove(filename)

    periods = 5
    index = pd.date_range('1/1/2016', periods=periods, freq='D')
    data = np.random.rand(periods, 4)
    df = pd.DataFrame(data=data, index=index, columns=['A', 'B', 'C', 'D'])

    plt.figure()
    pecos.graphics.plot_heatmap(df)
    plt.savefig(filename, format='png', bbox_inches='tight', pad_inches = 0)
    plt.close()

    assert_true(isfile(filename))