test_graphics.py 文件源码

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

    periods = 365*12
    index = pd.date_range('1/1/2016', periods=periods, freq='2H')
    data = np.random.rand(periods)
    df = pd.DataFrame(data=data, index=index, columns=['A'])
    overlay = pd.DataFrame(index=[1,100,200,300,365], 
                           data={'A': [40000,20000,60000,10000,5000],
                                 'B': [60000,70000,75000,50000,65000]})

    plt.figure()
    pecos.graphics.plot_doy_heatmap(df['A'], cmap='gray', overlay=overlay)
    plt.savefig(filename, format='png')
    plt.close()

    assert_true(isfile(filename))