def test_bollinger_bands(self):

        a = [10,12,12,13,9,12,12,13]
        q = pd.DataFrame(a)

        res=(pd.DataFrame([nan,nan,nan,nan,16.129503,16.149725,16.149725,16.729503]),
             pd.DataFrame([nan,nan,nan,nan,6.270497,7.050275,7.050275,6.870497]),
             pd.DataFrame([nan,nan,nan,nan,11.2,11.6,11.6,11.8]))

        assert_frame_equal(bollinger_bands(q,5,3)[0],res[0])
        assert_frame_equal(bollinger_bands(q,5,3)[1],res[1])
        assert_frame_equal(bollinger_bands(q,5,3)[2],res[2])
예제 #2
0
파일: aux.py 프로젝트: Omnitron/algotrading
def plot_data(data,
              name='',
              date=[0, 0],
              smas=var.default_smas,
              emas=var.default_emas,
              entry_points=None,
              exit_points=None,
              market_name='',
              to_file=False,
              show_smas=False,
              show_emas=False,
              show_bbands=False):
    '''
    Plots selected data.
    entry_points is a tuple of lists: (entry_points_x,entry_points_y)
    '''
    #plt.clf()

    # For when it's called outside backtest.
    if date != [0, 0]:
        if len(data) != date[1] - date[0]:
            data = data[date[0]:date[1]]

    f, (ax1, ax2, ax3) = plt.subplots(3,
                                      sharex=True,
                                      figsize=(9, 4),
                                      gridspec_kw={'height_ratios': [3, 1, 1]})

    ax1.grid(True)
    ax2.grid(True)
    ax3.grid(True)

    # var date is causing conflicts. using name date.
    if date[1] == 0:
        end_date = len(data)
    else:
        end_date = date[1]

    x = range(date[0], end_date)
    ax1.plot(x, data.Last, color='black', linewidth=1, alpha=0.65)

    if show_bbands:
        bb_upper, bb_lower, bb_sma = bollinger_bands(data.Last, 10, 2)
        #ax1.plot(x, bb_upper, color='red', linestyle='none', linewidth=1)
        #ax1.plot(x, bb_lower, color='green', linestyle='none', linewidth=1)

        ax1.fill_between(x, bb_sma, bb_upper, color='green', alpha=0.3)
        ax1.fill_between(x, bb_lower, bb_sma, color='red', alpha=0.3)
    if show_smas:
        for sma in smas:
            ax1.plot(x, data.Last.rolling(sma).mean())

    if show_emas:
        for ema in emas:
            ax1.plot(x, data.Last.ewm(ema).mean())

    if entry_points:
        ax1.plot(entry_points[0],
                 entry_points[1],
                 marker='o',
                 linestyle='None',
                 color='green',
                 alpha=0.75)
    if exit_points:
        ax1.plot(exit_points[0],
                 exit_points[1],
                 marker='o',
                 linestyle='None',
                 color='red',
                 alpha=0.75)
    ax2.bar(x, data.BaseVolume.iloc[:], 1, color='black', alpha=0.55)

    try:
        ax3.plot(x, data.OpenSell.iloc[:])
    except Exception as e:
        ax3.plot(x, data.High.iloc[:])

    plt.xlim(date[0], end_date)
    plt.tight_layout()
    f.subplots_adjust(hspace=0)
    if to_file:
        if not name:
            name = 'fig_test' + str(time())
        f.savefig(var.fig_dir + name + '.pdf', bbox_inches='tight')
        plt.close(f)
    #plt.show()

    return True