简单移动平均线(SMA)的功能

时间:2017-03-02 11:13:48

标签: python

我有以下函数用于在python中计算SMA:

import numpy as np

def calcSma(data, smaPeriod):
    sma = []
    count = 0
    for i in xrange(data.size):
        if data[i] is None:
            sma.append(None)
        else:
            count += 1
            if count < smaPeriod:
                sma.append(None)
            else:
                sma.append(np.mean(data[i-smaPeriod+1:i+1]))

    return np.array(sma)

这个功能有效,但我发现它很少pythonic。我不喜欢我正在做的索引和计数,也不喜欢我必须附加到列表然后在返回它之前把它变成一个numpy数组的方式。

我必须处理所有这些None的原因是因为我想返回一个与输入数组大小相同的数组。这使得以后更容易绘制和处理一般级别。我可以很容易地做到这样的事情:

sma = calcSma(data=data, smaPeriod=20)
sma2 = calcSma(data=sma, smaPeriod=10)
plt.plot(data)
plt.plot(sma)
plt.plot(sma2)
plt.show()

那么,关于如何做到这一点更漂亮和更加pythonic的任何想法?

2 个答案:

答案 0 :(得分:0)

Pythonic我希望

import numpy as np


def calcSma(data, smaPeriod):
    j = next(i for i, x in enumerate(data) if x is not None)
    our_range = range(len(data))[j + smaPeriod - 1:]
    empty_list = [None] * (j + smaPeriod - 1)
    sub_result = [np.mean(data[i - smaPeriod + 1: i + 1]) for i in our_range]

    return np.array(empty_list + sub_result)

答案 1 :(得分:-2)

这里是仅使用标准Python库进行移动平均的另一种实现方式:

from collections import deque
import itertools

def moving_average(iterable, n=3):
    # http://en.wikipedia.org/wiki/Moving_average
    it = iter(iterable) 
    # create an iterable object from input argument
    d = deque(itertools.islice(it, n-1))  
    # create deque object by slicing iterable
    d.appendleft(0)
    s = sum(d)
    for elem in it:
        s += elem - d.popleft()
        d.append(elem)
        yield s / n

# example on how to use it
for i in  moving_average([40, 30, 50, 46, 39, 44]):
    print(i)

# 40.0
# 42.0
# 45.0
# 43.0