1
$\begingroup$

Usually called WMA. The weighting is linear (as opposed to exponential) defined here: Moving Average, Weighted. I attempt to implement this in a python function as show below. The result is a list of values. My question is: are the result right? Also it is very slow...

I input a dataframe from pandas with a column called 'close'

def wma(df):
    n = 20
    k = (n * (n + 1)) / 2.0
    wmas = []
    for i in range(0, len(df) - n + 1):
        product = [df['close'][i + n_i] * (n_i + 1) for n_i in range(0, n)]
        wma = sum(product) / k
        wmas.append(wma)
    return wmas

Any help would be appreciated. Thanks.

$\endgroup$
2
  • 1
    $\begingroup$ The code looks like it will give the correct result but it is not dynamic and you are asking for errors being that you have the period length and column name hardcoded. They would typically be passed variables. There are much faster implementations of this without using a for loop with and nested list comprehension. This is almost painful to look at! $\endgroup$
    – amdopt
    Sep 1 '20 at 19:26
  • $\begingroup$ @amdopt any advice would be help. My code takes almost 5 seconds with a column of just 5,000 values! $\endgroup$
    – user89135
    Sep 1 '20 at 19:29
2
$\begingroup$

Though your code is already giving you the correct result, I almost feel bad for you that you have to wait 5 seconds for such a small amount of data. Your code is slow because you are kind of reinventing the wheel instead of using some built-in pandas and numpy functionality. For example, product and wma in your code can be combined and accomplished using numpy's dot product function (np.dot) that is applied to the whole column in a rolling fashion with an anonymous function by chaining pandas .rolling() and .apply() methods. It is always better to look for ready-made solutions becuase the functions are optimized behind the scenes. I ran your code on my machine, and the results take about 2 seconds for 5200 values. Try something like this (I added some basic functionality as an example to get you thinking):

import pandas as pd
import numpy as np

def wma(df, column='close', n=20, add_col=False):

    weights = np.arange(1, n + 1)
    wmas = df[column].rolling(n).apply(lambda x: np.dot(x, weights) /
                                       weights.sum(), raw=True).to_list()

    if add_col == True:
        df[f'{column}_WMA_{n}'] = wmas
        return df
    else:
        return wmas

The above function will take the same dataframe you were using and return the same list the same way you had it--just call wma(df). In addition can change the column name, the period value, and you can opt to not return a list but add the values as a new column to the dataframe that you originally passed in. It also runs on my machine in about 20 milliseconds--almost 100x faster than your original code...

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.