计算布林带

# coding: utf-8
# 作者：Wizard <github.com/wizardforcel>

import numpy as np
from matplotlib import pyplot as plt
import sys

# 获取数据
# 因为没找到数据源，所以直接随机生成了
series = np.random.rand(100) * 10 + 10
ser_len = len(series)

# 获取窗口大小，默认为 5
win_sz = 5 if len(sys.argv) < 2 else sys.argv[1]
print u'窗口大小: ', win_sz

# 计算权重，sma_weight 为简单滑动
# ema_weight 为指数滑动
sma_weight = np.ones(win_sz).astype(float) / win_sz
ema_weight = np.exp(np.linspace(-1, 0, win_sz))
ema_weight /= ema_weight.sum()


# 手动计算简单/指数滑动平均，以及滑动标准差
# 滑动平均就是当天和前 N - 1 天的（加权）均值
# 其中 N 是窗口大小，标准差与之类似
sma = []
ema = []
mstd = []

for i in xrange(win_sz - 1, ser_len):

    ser_range = series[i - (win_sz - 1): i + 1]
    avg = (ser_range * sma_weight).sum()
    sma.append(avg)
    avg = (ser_range * ema_weight).sum()
    ema.append(avg)

    std = ser_range.std()
    mstd.append(std)

sma = np.asarray(sma)
mstd = np.asarray(mstd)

# 上布林带是简单滑动均值加上两个滑动标准差
# 下布林带是简单滑动均值减去两个滑动标准差
upper = sma + 2 * mstd
lower = sma - 2 * mstd

# 要注意横轴从 N - 1 开始
ser_slice = series[win_sz - 1:]
x = np.arange(win_sz - 1, ser_len)
print 'x: ', x[:5]

# 绘制函数
plt.plot(x, ser_slice, 'b', label='series')
print u'序列: ', ser_slice[:5]
plt.plot(x, sma, 'g', label='sma')
print u'简单滑动均值: ', sma[:5]
plt.plot(x, ema, 'r', label='ema')
print u'指数滑动均值: ', ema[:5]

plt.plot(x, upper, 'grey', label='upper bolling')
plt.plot(x, lower, 'grey', label='lower bolling')

plt.legend()
plt.show()