2

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?

Panto
  • 104
  • 1
  • 1
  • 9
  • 1
    Possible duplicate of [Moving average or running mean](http://stackoverflow.com/questions/13728392/moving-average-or-running-mean) – Dschoni Mar 02 '17 at 11:17
  • This makes your code crash: `calcSma(np.array([1, 1, 1, 1, 1, 1, None, 1, 1, 1]), 3)`, can you confirm? – Elmex80s Mar 02 '17 at 13:33
  • Yes, it will crash, but since the None elements always will be in the start of an array that's not really a problem. – Panto Mar 02 '17 at 17:29
  • Here is my solution just using standard Python library: [Moving average or running mean](https://stackoverflow.com/a/52273625/30038) – Vlad Bezden Sep 11 '18 at 10:19
  • 1
    I just used my data frames' rolling function: `myDataFrame["new_sma_col"] = myDataFrame["my_col"].rolling(10).mean()` –  Apr 19 '21 at 19:09

2 Answers2

4

Pythonic enough I hope

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)
Elmex80s
  • 3,428
  • 1
  • 15
  • 23
  • 1
    Well, the list comprehension helps, but without appending None it don't work the way I want. – Panto Mar 02 '17 at 12:56
  • I did a small update. Can you confirm this is how it should work? – Elmex80s Mar 02 '17 at 13:46
  • Almost, but it still fails when the data array contains None elements. As in my first example, where I put sma into the function to calculate sma2. Also, in numpy we need to do np.append(array1, array2) instead of array1 + array2. – Panto Mar 02 '17 at 17:45
  • Yes, it works now. However, I'm not sure if it turned out more pythonic or readable. It's been helpful anyway, thank you! – Panto Mar 07 '17 at 10:46
-2

Here is another implementation of moving average just using standard Python library:

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
Vlad Bezden
  • 83,883
  • 25
  • 248
  • 179