correlation.py 文件源码

def plot_index_and_sentiment(tick_seq, shindex_seq, sentiment_seq, date):
    if len(tick_seq) != len(shindex_seq) or len(tick_seq) != len(sentiment_seq):
        print('error(plot) : three sequence length is not same')
        return

    x = range(len(shindex_seq))
    labels = tick_seq
    y1 = shindex_seq
    y2 = sentiment_seq

    def format_fn(tick_val, tick_pos):
        if int(tick_val) in x:
            return labels[int(tick_val)]
        else:
            return ''

    fig = plt.figure(figsize=(12,8))
    p1 = fig.add_subplot(111)
    p1.xaxis.set_major_formatter(FuncFormatter(format_fn))
    p1.xaxis.set_major_locator(MaxNLocator(integer=True, nbins=12))
    delta = shindex_seq[len(shindex_seq) - 1] - shindex_seq[0]
    if delta > 0:
        p1.plot(x, y1, label="$SCI$", color="red", linewidth=1)
    else:
        p1.plot(x, y1, label="$SCI$", color="green", linewidth=1)
    p1.plot(x, y2, 'b--', label="$ISI$", color="blue", linewidth=1)

    plt.title("Shanghai Composite Index(SCI) & Investor Sentiment Index(ISI)")
    plt.xlabel("Time(5min)")
    plt.ylabel("Index Value")
    plt.legend()
    # plt.show()
    global subdir
    filepath = './Pic/' + subdir + '/' + date + '.png'
    plt.savefig(filepath)