Python绘制加强版K线图：增加均线及成交量

在之前，我们讲解了如何用Python绘制K线图。当时就有人问能否加入均线元素，当然能啦！不光均线，今天我还要把成交量给加上去。

获取数据并加工

为了计算年线（250日均线），我们需要保证数据至少在一年以上，所以我们获取17年以来上证综指的行情数据。

获取之后，我们同时保留两种格式的日期数据：matplotlib需要的float型以及我们坐标轴标签中需要的字符串型。它们分别是下边的trade_date列和trade_date2列。

另外我们需要一串连续的数字（最好是从零开始可以作为索引的），用来保证我们的K线图是连续的，也就是下边的dates列。

对了，别忘了对数据按照日期升序排列，不然得到的可能就是从最新日期开始计算的移动均值了。

import mpl_finance
import tushare as ts
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib import ticker
from matplotlib.pylab import date2num
import numpy as np
import matplotlib as mpl
sns.set()
mpl.rcParams['font.family'] = 'sans-serif'
mpl.rcParams['font.sans-serif'] = 'SimHei'
pro = ts.pro_api()

df = pro.index_daily(ts_code='000001.SH', start_date='20170101')
df = df.sort_values(by='trade_date', ascending=True)
df['trade_date2'] = df['trade_date'].copy()
df['trade_date'] = pd.to_datetime(df['trade_date']).map(date2num)
df['dates'] = np.arange(0, len(df))
df.head()

用窗口函数计算移动平均值

df['5'] = df.close.rolling(5).mean()
df['20'] = df.close.rolling(20).mean()
df['30'] = df.close.rolling(30).mean()
df['60'] = df.close.rolling(60).mean()
df['120'] = df.close.rolling(120).mean()
df['250'] = df.close.rolling(250).mean()

绘制K线图

def format_date(x,pos):
    if x<0 or x>len(date_tickers)-1:
        return ''
    return date_tickers[int(x)]

date_tickers = df.trade_date2.values
fig, ax = plt.subplots(figsize=(10,5))
ax.xaxis.set_major_formatter(ticker.FuncFormatter(format_date))
# 绘制K线图
mpl_finance.candlestick_ochl(
    ax=ax,
    quotes=df[['dates', 'open', 'close', 'high', 'low']].values,
    width=0.7,
    colorup='r',
    colordown='g',
    alpha=0.7)
# 绘制均线
for ma in ['5', '20', '30', '60', '120', '250']:
    plt.plot(df['dates'], df[ma])
plt.legend()
ax.set_title('上证综指K线图(2017.1-)', fontsize=20);

由于战线太长，显示效果不是特别好，我们缩短到最近半年左右：

df2 = df.query('trade_date2 >= "20180601"').reset_index()
df2['dates'] = np.arange(0, len(df2))
date_tickers = df2.trade_date2.values

fig, ax = plt.subplots(figsize=(10,5))
ax.xaxis.set_major_formatter(ticker.FuncFormatter(format_date))

# 绘制K线图
mpl_finance.candlestick_ochl(
    ax=ax,
    quotes=df2[['dates', 'open', 'close', 'high', 'low']].values,
    width=0.7,
    colorup='r',
    colordown='g',
    alpha=0.7)

# 绘制均线
for ma in ['5', '20', '30', '60', '120', '250']:
    plt.plot(df2['dates'], df2[ma])
plt.legend()
ax.set_title('上证综指K线图(2018.6-)', fontsize=20);

看起来清晰了很多。可以看到，短期30日线压力明显，中长期60日线压力明显。一旦实现有效突破，就算没个大行情，这难熬的熊市也基本上可以宣告尾声了。

绘制成交量

你以为这就结束了吗？哈哈，你太天真了。我说过有成交量的，那就肯定会有成交量。

我们需要用到matplotlib.gridspec来控制子图的位置和大小；然后针对ax1绘制K线图，针对ax2绘制成交量。由于阳线为红色，阴线为绿色，所以我们把成交量的颜色也对应着调整下。

from matplotlib.gridspec import GridSpec

# 取18.9以来数据
df2 = df.query('trade_date2 >= "20180601"').reset_index()
df2['dates'] = np.arange(0, len(df2))
date_tickers = df2.trade_date2.values

# 控制子图
figure = plt.figure(figsize=(12, 9))
gs = GridSpec(3, 1)
ax1 = plt.subplot(gs[:2, :])
ax2 = plt.subplot(gs[2, :])

# 绘制K线图
mpl_finance.candlestick_ochl(
    ax=ax1,
    quotes=df2[['dates', 'open', 'close', 'high', 'low']].values,
    width=0.7,
    colorup='r',
    colordown='g',
    alpha=0.7)
ax1.xaxis.set_major_formatter(ticker.FuncFormatter(format_date))

# 绘制均线
for ma in ['5', '20', '30', '60', '120', '250']:
    ax1.plot(df2['dates'], df2[ma])
ax1.legend()
ax1.set_title('上证综指K线图(2018.6-)', fontsize=20)
ax1.set_ylabel('指数')

# 绘制成交量
ax2.xaxis.set_major_formatter(ticker.FuncFormatter(format_date))
df2['up'] = df.apply(lambda row: 1 if row['close'] >= row['open'] else 0, axis=1)
ax2.bar(df2.query('up == 1')['dates'], df2.query('up == 1')['vol'], color='r', alpha=0.7)
ax2.bar(df2.query('up == 0')['dates'], df2.query('up == 0')['vol'], color='g', alpha=0.7)
ax2.set_ylabel('成交量')
#plt.xticks(date_tickers);

这样看起来不真切，我们调整到9月份开始来看看：

好了，K线图的绘制就暂时告一段落。之后如果需要整合各种指标的话，我们再来绘制一个更加强大的版本，敬请期待哦！