一步一步教你用Python画出专业的K线图

由polll创建，最终由polll更新于2023-02-13 08:42 被浏览 150 用户

漂亮的界面是量化程序的脸面，直观专业的界面能帮助我们了解股票的走势和形象地展示量化交易的结果。本文介绍如何使用matplotlib做出专业的K线图和展示各种技术指标：


matplotlib是受MATLAB的启发构建的。MATLAB是数据绘图领域广泛使用的语言和工具。MATLAB语言是面向过程的。利用函数的调用，MATLAB中可以轻松的利用一行命令来绘制直线，然后再用一系列的函数调整结果。

matplotlib有一套完全仿照MATLAB的函数形式的绘图接口，十分方便实用。下面我们一步步使用它来绘出上面的图片。

准备工作

在上文中（量化回测的基础：干净的，准确的历史股票数据），我们已经为大家准备好了日k线数据，大家可以去下载。下面我们来为绘图做一些准备工作：

from pandas import DataFrame, Series
import pandas as pd; import numpy as np
import matplotlib.pyplot as plt
from matplotlib import dates as mdates
from matplotlib import ticker as mticker
from matplotlib.finance import candlestick_ohlc
from matplotlib.dates import DateFormatter, WeekdayLocator, DayLocator, MONDAY,YEARLY
from matplotlib.dates import MonthLocator,MONTHLY
import datetime as dt
import pylab

daylinefilespath = 'G:\\dayline\\'
stock_b_code = '000001' #平安银行
MA1 = 10
MA2 = 50
startdate = dt.date(2016, 6, 29)
enddate = dt.date(2017, 1, 30)


def readstkData(rootpath, stockcode, sday, eday):
    
    returndata = pd.DataFrame()
    for yearnum in range(0,int((eday - sday).days / 365.25)+1):
        theyear = sday + dt.timedelta(days = yearnum * 365)
        # build file name
        filename = rootpath  + theyear.strftime('%Y') + '\\' + str(stockcode).zfill(6) + '.csv'
        
        try:
            rawdata = pd.read_csv(filename, parse_dates = True, index_col = 0, encoding = 'gbk')
        except IOError:
           raise Exception('IoError when reading dayline data file: ' + filename)

        returndata = pd.concat([rawdata, returndata])
    
    # Wash data
    returndata = returndata.sort_index()
    returndata.index.name = 'DateTime'
    returndata.drop('amount', axis=1, inplace = True)
    returndata.columns = ['Open', 'High', 'Close', 'Low', 'Volume']

    returndata = returndata[returndata.index < eday.strftime('%Y-%m-%d')]

    return returndata

这里定义了几个全局变量：

1. daylinefilespath 是下载的日k线数据
2. stock_b_code 是要显示的股票代码，大家可以换成自己想要的股票
3. MA1和MA2是移动平均线的日期间隔，例如10是10日移动平均线
4. startdate和enddate是想要显示的日期范围

函数readstkData读入数据并做了必要的合并和清洗。注意这里年做了简化处理，并不完全准确。

绘出日K线

Talk is cheap，我们直接上代码：

def main():
    days = readstkData(daylinefilespath, stock_b_code, startdate, enddate)

    # drop the date index from the dateframe & make a copy
    daysreshape = days.reset_index()
    # convert the datetime64 column in the dataframe to 'float days'
    daysreshape['DateTime']=mdates.date2num(daysreshape['DateTime'].astype(dt.date))
    # clean day data for candle view 
    daysreshape.drop('Volume', axis=1, inplace = True)
    daysreshape = daysreshape.reindex(columns=['DateTime','Open','High','Low','Close'])  
    
    Av1 = movingaverage(daysreshape.Close.values, MA1)
    Av2 = movingaverage(daysreshape.Close.values, MA2)
    SP = len(daysreshape.DateTime.values[MA2-1:])
    fig = plt.figure(facecolor='#07000d',figsize=(15,10))
    
    ax1 = plt.subplot2grid((6,4), (1,0), rowspan=4, colspan=4, axisbg='#07000d')
    candlestick_ohlc(ax1, daysreshape.values[-SP:], width=.6, colorup='#ff1717', colordown='#53c156')
    Label1 = str(MA1)+' SMA'
    Label2 = str(MA2)+' SMA'
    
    ax1.plot(daysreshape.DateTime.values[-SP:],Av1[-SP:],'#e1edf9',label=Label1, linewidth=1.5)
    ax1.plot(daysreshape.DateTime.values[-SP:],Av2[-SP:],'#4ee6fd',label=Label2, linewidth=1.5)
    ax1.grid(True, color='w')
    ax1.xaxis.set_major_locator(mticker.MaxNLocator(10))
    ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
    ax1.yaxis.label.set_color("w")
    ax1.spines['bottom'].set_color("#5998ff")
    ax1.spines['top'].set_color("#5998ff")
    ax1.spines['left'].set_color("#5998ff")
    ax1.spines['right'].set_color("#5998ff")
    ax1.tick_params(axis='y', colors='w')
    plt.gca().yaxis.set_major_locator(mticker.MaxNLocator(prune='upper'))
    ax1.tick_params(axis='x', colors='w')
    plt.ylabel('Stock price and Volume')
    plt.show()

if __name__ == "__main__":
    main()

这里有几点要注意：

1. 为了调用matplotlib的蜡烛图函数，我们把数据重新排布了一下，并修改了日期的格式。以满足candlestick_ohlc函数的要求。
2. 创立了ax1来显示日k线，后面还会建立其他几个子图。

我们来运行一下：


效果还行，不但有K线，还叠加了MA1和MA2。下面我们来绘制其他技术参数。

绘制RSI

下面我们在顶部加入RSI:

 # plot an RSI indicator on top
    maLeg = plt.legend(loc=9, ncol=2, prop={'size':7},
                       fancybox=True, borderaxespad=0.)
    maLeg.get_frame().set_alpha(0.4)
    textEd = pylab.gca().get_legend().get_texts()
    pylab.setp(textEd[0:5], color = 'w')
    
    ax0 = plt.subplot2grid((6,4), (0,0), sharex=ax1, rowspan=1, colspan=4, axisbg='#07000d')
    rsi = rsiFunc(daysreshape.Close.values)
    rsiCol = '#c1f9f7'
    posCol = '#386d13'
    negCol = '#8f2020'
    
    ax0.plot(daysreshape.DateTime.values[-SP:], rsi[-SP:], rsiCol, linewidth=1.5)
    ax0.axhline(70, color=negCol)
    ax0.axhline(30, color=posCol)
    ax0.fill_between(daysreshape.DateTime.values[-SP:], rsi[-SP:], 70, where=(rsi[-SP:]>=70), facecolor=negCol, edgecolor=negCol, alpha=0.5)
    ax0.fill_between(daysreshape.DateTime.values[-SP:], rsi[-SP:], 30, where=(rsi[-SP:]<=30), facecolor=posCol, edgecolor=posCol, alpha=0.5)
    ax0.set_yticks([30,70])
    ax0.yaxis.label.set_color("w")
    ax0.spines['bottom'].set_color("#5998ff")
    ax0.spines['top'].set_color("#5998ff")
    ax0.spines['left'].set_color("#5998ff")
    ax0.spines['right'].set_color("#5998ff")
    ax0.tick_params(axis='y', colors='w')
    ax0.tick_params(axis='x', colors='w')
    plt.ylabel('RSI')

我们在ax1的上面加上了ax0子图，用以绘制RSI曲线。程序比较简单，我们来运行一下：


绘制成交量

所有的股票软件都有成交量图，我们当然也不能免俗。我们不再增加子图，而是把成交量叠加在k线图下面，用浅蓝色标识，这样既清楚又节省空间：

 volumeMin = 0
    ax1v = ax1.twinx()
    ax1v.fill_between(daysreshape.DateTime.values[-SP:],volumeMin, days.Volume.values[-SP:], facecolor='#00ffe8', alpha=.4)
    ax1v.axes.yaxis.set_ticklabels([])
    ax1v.grid(False)
    ###Edit this to 3, so it's a bit larger
    ax1v.set_ylim(0, 3*days.Volume.values.max())
    ax1v.spines['bottom'].set_color("#5998ff")
    ax1v.spines['top'].set_color("#5998ff")
    ax1v.spines['left'].set_color("#5998ff")
    ax1v.spines['right'].set_color("#5998ff")
    ax1v.tick_params(axis='x', colors='w')
    ax1v.tick_params(axis='y', colors='w')

注意这里不是家里一个新的ax，而是从ax1上加入了个twinx，从而和k线共用x轴。效果如下：


绘制MACD

我们再在最下面画出MACD图。经过了前面那些步骤，你一定知道只要增加个新的子图就好了。是的，程序比较简单：

 # plot an MACD indicator on bottom
    ax2 = plt.subplot2grid((6,4), (5,0), sharex=ax1, rowspan=1, colspan=4, axisbg='#07000d')
    fillcolor = '#00ffe8'
    nslow = 26
    nfast = 12
    nema = 9
    emaslow, emafast, macd = computeMACD(daysreshape.Close.values)
    ema9 = ExpMovingAverage(macd, nema)
    ax2.plot(daysreshape.DateTime.values[-SP:], macd[-SP:], color='#4ee6fd', lw=2)
    ax2.plot(daysreshape.DateTime.values[-SP:], ema9[-SP:], color='#e1edf9', lw=1)
    ax2.fill_between(daysreshape.DateTime.values[-SP:], macd[-SP:]-ema9[-SP:], 0, alpha=0.5, facecolor=fillcolor, edgecolor=fillcolor)
     plt.gca().yaxis.set_major_locator(mticker.MaxNLocator(prune='upper'))
    ax2.spines['bottom'].set_color("#5998ff")
    ax2.spines['top'].set_color("#5998ff")
    ax2.spines['left'].set_color("#5998ff")
    ax2.spines['right'].set_color("#5998ff")
    ax2.tick_params(axis='x', colors='w')
    ax2.tick_params(axis='y', colors='w')
    plt.ylabel('MACD', color='w')
    ax2.yaxis.set_major_locator(mticker.MaxNLocator(nbins=5, prune='upper'))
    for label in ax2.xaxis.get_ticklabels():
        label.set_rotation(45)   

运行效果如下：


收尾

不知你注意到没有，上图有几个问题：

1. x轴日期出现了三次
2. 没有title

下面我们做一些美化工作，并演示如何加入一个提示点，将来你可以用它来展示自己量化操作的节点：

 plt.suptitle(stock_b_code,color='w')
    plt.setp(ax0.get_xticklabels(), visible=False)
    plt.setp(ax1.get_xticklabels(), visible=False)
    
    # Mark big event
    # TODO: Make a real case here
    ax1.annotate('BreakNews!',(daysreshape.DateTime.values[155],Av1[155]),
        xytext=(0.8, 0.9), textcoords='axes fraction',
        arrowprops=dict(facecolor='white', shrink=0.05),
        fontsize=10, color = 'w',
        horizontalalignment='right', verticalalignment='bottom')
    
    plt.subplots_adjust(left=.09, bottom=.14, right=.94, top=.95, wspace=.20, hspace=0)
    plt.show()

好了，大功告成：


结语

大家可能注意到了，有很多代码在配色。是的，如果我们把所有的colorxxx都删掉，也可以显示出大致的曲线，但是却非常难看。这也是为什么要改变缺省matplotlib配色的原因。

大家可以实验一下更新前面的几个全局变量，比如换一换股票、MA或者开始结束日期什么的，这里给大家一个思考题：

在readstkData（）切片时为什么只切了结束时间没有切开始时间？SP变量是干什么用的？

欢迎关注本专栏！