[021量化交易] python 绘画K线 MACD RSI 成交量 指标

import talib
import matplotlib.pyplot as plt
from mplfinance.original_flavor import candlestick_ohlc
import matplotlib.ticker as ticker

# 导入tushare
import tushare as ts

# 初始化pro接口
pro = ts.pro_api('')

# 拉取数据
data = pro.daily(**{
    "ts_code": "000001.SZ",
    "trade_date": "",
    "start_date": 20220101,
    "end_date": "20220601",
    "offset": "",
    "limit": ""
}, fields=[
    "trade_date",
    "open",
    "high",
    "low",
    "close",
    "vol",
    "pre_close"
])

# df = df.reset_index()  # 默认index是日期，这里要重置一下，为后面绘图做准备
# df = df.astype({"trade_date": str})  # 将Date列的类型设置为str，为绘图做准备

# print(df)


# 获取macd
data["macd"], data["macd_signal"], data["macd_hist"] = talib.MACD(data['close'])

# 获取10日均线和30日均线
data["ma7"] = talib.MA(data["close"], timeperiod=7)
data["ma8"] = talib.MA(data["close"], timeperiod=8)
data["ma10"] = talib.MA(data["close"], timeperiod=10)
data["ma25"] = talib.MA(data["close"], timeperiod=25)
data["ma30"] = talib.MA(data["close"], timeperiod=30)

# 获取rsi
data["rsi"] = talib.RSI(data["close"])

print(data)

fig = plt.figure()  # 创建绘图区，包含四个子图
fig.set_size_inches((20, 16))
ax_candle = fig.add_axes((0, 0.72, 1, 0.32))  # 蜡烛图子图
ax_macd = fig.add_axes((0, 0.48, 1, 0.2), sharex=ax_candle)  # macd子图
ax_rsi = fig.add_axes((0, 0.24, 1, 0.2), sharex=ax_candle)  # rsi子图
ax_vol = fig.add_axes((0, 0, 1, 0.2), sharex=ax_candle)  # 成交量子图

ohlc = []  # 存放行情数据，candlestick_ohlc需要传入固定格式的数据
row_number = 0
for date, row in data.iterrows():
    date, openp, highp, lowp, closep = row[:5]
    ohlc.append([row_number, openp, highp, lowp, closep])
    row_number = row_number + 1

date_tickers = data.trade_date.values  # 获取Date数据


def format_date(x, pos=None):
    # 由于前面股票数据在 date 这个位置传入的都是int
    # 因此 x=0,1,2,...
    # date_tickers 是所有日期的字符串形式列表
    if x < 0 or x > len(date_tickers) - 1:
        return ''
    return date_tickers[int(x)]


# 绘制蜡烛图
ax_candle.plot(data.index, data["ma7"], label="MA7")
ax_candle.plot(data.index, data["ma8"], label="MA8")
ax_candle.plot(data.index, data["ma25"], label="MA25")
candlestick_ohlc(ax_candle, ohlc, colorup="g", colordown="r", width=0.8)

ax_candle.xaxis.set_major_formatter(ticker.FuncFormatter(format_date))
ax_candle.xaxis.set_major_locator(ticker.MultipleLocator(6))  # 设置间隔为6个交易日
ax_candle.grid(True)
ax_candle.set_title("title", fontsize=20)
ax_candle.legend()

# 绘制MACD
ax_macd.plot(data.index, data["macd"], label="macd")
ax_macd.bar(data.index, data["macd_hist"] * 3, label="hist")
ax_macd.plot(data.index, data["macd_signal"], label="signal")
ax_macd.set_title('MACD')
ax_macd.legend()

# 绘制RSI
ax_rsi.set_ylabel("(%)")
ax_rsi.plot(data.index, [70] * len(data.index), label="overbought")
ax_rsi.plot(data.index, [30] * len(data.index), label="oversold")
ax_rsi.plot(data.index, data["rsi"], label="rsi")
ax_rsi.set_title('KDJ')
ax_rsi.legend()

# 绘制成交量
ax_vol.bar(data.index, data["vol"] / 1000000)
ax_vol.set_ylabel("(Million)")

# 保存图片到本地
# fig.savefig("/Users/answer/Desktop/investment/photos/" + title + ".png", bbox_inches="tight")

# 这里个人选择不要plt.show()，因为保存图片到本地的
plt.show()