Calculando uma média móvel linear ponderada em Python

Question

Normalmente chamado de WMA. A ponderação é linear (em oposição a exponencial) definida aqui: Média móvel, ponderada . Tento implementar isso em uma função python, conforme mostrado abaixo. O resultado é uma lista de valores. Minha pergunta é: o resultado está certo? Também é muito lento ...

Eu insiro um dataframe de pandas com uma coluna chamada 'fechar'

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

Qualquer ajuda seria apreciada. Obrigado.

Answer 1

Embora seu código já esteja fornecendo o resultado correto, quase me sinto mal por você ter que esperar 5 segundos por uma quantidade tão pequena de dados. Seu código é lento porque você está reinventando a roda em vez de usar alguns pandas embutidos e funcionalidades entorpecidas. Por exemplo, producte wmaem seu código podem ser combinados e realizados usando a função de produto escalar de numpy ( np.dot) que é aplicada a toda a coluna de forma contínua com uma função anônima, encadeando pandas .rolling()e.apply()métodos. É sempre melhor procurar soluções prontas porque as funções são otimizadas nos bastidores. Executei seu código na minha máquina e os resultados demoram cerca de 2 segundos para 5200 valores. Tente algo assim (adicionei algumas funcionalidades básicas como exemplo para fazer você pensar):

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

A função acima pegará o mesmo dataframe que você estava usando e retornará a mesma lista da mesma forma que você - basta chamar wma(df). Além disso, pode alterar o nome da coluna, o valor do período e você pode optar por não retornar uma lista, mas adicionar os valores como uma nova coluna ao dataframe que você transmitiu originalmente. Ele também é executado em minha máquina em cerca de 20 milissegundos - quase 100x mais rápido que seu código original ...