Calculando un Promedio Móvil Ponderado Lineal en Python

Question

Generalmente llamado WMA. El peso es lineal (a diferencia del exponencial) definido aquí: Media Móvil, Ponderada. Intento implementar esto en una función de python como se muestra a continuación. El resultado es una lista de valores. ¿Son los resultados correctos? Además, es muy lento...

Inserto un dataframe de pandas con una columna llamada 'close'

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

Cualquier ayuda sería apreciada. Gracias.

Answer 1

Aunque tu código ya te está dando el resultado correcto, casi me siento mal por ti porque tienes que esperar 5 segundos por una cantidad tan pequeña de datos. Tu código es lento porque estás reinventando la rueda en lugar de utilizar algunas funciones integradas de pandas y numpy. Por ejemplo, producto y wma en tu código pueden combinarse y lograrse utilizando la función de producto punto de numpy (np.dot) que se aplica a toda la columna de forma iterativa con una función anónima encadenando los métodos .rolling() y .apply() de pandas. Siempre es mejor buscar soluciones predefinidas porque las funciones están optimizadas detrás de escena. Ejecuté tu código en mi máquina y los resultados tardan alrededor de 2 segundos para 5200 valores. Prueba algo así (agregué una funcionalidad básica como ejemplo para que empieces a 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

La función anterior tomará el mismo dataframe que estabas usando y devolverá la misma lista de la misma manera que la tenías originalmente. Solo llama a wma(df). Además, puedes cambiar el nombre de la columna, el valor del período, y puedes optar por no devolver una lista, sino agregar los valores como una nueva columna al dataframe que pasaste originalmente. Además, se ejecuta en mi máquina en aproximadamente 20 milisegundos, ¡casi 100 veces más rápido que tu código original...