Exponential moving average à la pandas

ewm.py

import numpy as np


def ewma(x, alpha, adjust=True):
    q = 1 - alpha
    y = x.copy()

    if not adjust:
        for i in range(1, len(y)):
            y[i] += q * (y[i-1] - y[i])
        return y

    for i in range(1, len(y)):
        y[i] += q * y[i-1]

    n = np.arange(len(x))
    s = (1 - q**(n+1)) / (1 - q)
    return y / s


def ewmcov(x, y, alpha, bias=False, adjust=True):
    if x is y:
        z = ewma(x**2, alpha, adjust=adjust)
        z -= ewma(x, alpha, adjust=adjust)**2
    else:
        z = ewma(x * y, alpha, adjust=adjust)
        z -= ewma(x, alpha, adjust=adjust) * ewma(y, alpha, adjust=adjust)

    if bias:
        return z

    q = 1 - alpha
    n = np.arange(len(x))
    if adjust:
        s2 = (1 - (q**2)**(n+1)) / (1 - q**2)
        s1 = (1 - q**(n+1)) / (1 - q)
    else:
        b = alpha**2 / (1 - q**2)
        s2 = (1 - b) * q**(2*n) + b
        s1 = 1

    bias = 1 - s2 / s1**2
    bias[0] = np.NaN
    return z / bias


def ewmvar(x, alpha, bias=False, adjust=True):
    return ewmcov(x, x, alpha, bias=bias, adjust=adjust)


def ewmstd(x, alpha, bias=False, adjust=True):
    return np.sqrt(ewmvar(x, alpha, bias=bias, adjust=adjust))


def ewmcorr(x, y, alpha, adjust=True):
    n = ewmcov(x, y, alpha, bias=True, adjust=adjust)
    d = ewmstd(x, alpha, bias=True, adjust=adjust)
    d *= ewmstd(y, alpha, bias=True, adjust=adjust)
    return n / d