I have the following function for calculating SMA in python:
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)
This function works, but I find it very little pythonic. I don't like the indexing and counting I'm doing, nor the way I have to append to the list and then turn it into a numpy array before I return it.
The reason I have to deal with all these None, is because I want to return an array at the same size as the input array. This makes it easier to plot and deal with on a general level later. I can easily do stuff such as this:
sma = calcSma(data=data, smaPeriod=20)
sma2 = calcSma(data=sma, smaPeriod=10)
plt.plot(data)
plt.plot(sma)
plt.plot(sma2)
plt.show()
So, any ideas on how this can be done prettier and more pythonic?