1
\$\begingroup\$

The following numpy/python function computes exponentially-weighted moving mean and standard deviation of an irregularly-spaced weighted time series. I want to make it faster by getting rid of the python loop and relying on numpy vectorization. However, while I can move some computation out of the loop, I'm unable to come up with a way to untangle the dependencies that exist between consecutive mean and variance values. Is it even possible?

I'd rather avoid Numba, Cython or the like.

import numpy as np


def ewm(values, time_steps, weights, decay_rate=0.01, variance_epsilon=1e-6):
    '''Compute exponentially-weighted moving mean and standard deviation of an
    irregularly-spaced weighted time series.
       
    :param values: 1D array of event values.
    :param time_steps: 1D array of floats. Each value indicates how much time
        has passed since the previous event (or the start of
        the series for the first event).
    :param weights: 1D array of floats. Each value indicates how much an event
        contributes to the moving mean and standard deviation.
    :param decay_rate: float. Describes how much influence of past events
        declines per unit of time.
    :param variance_epsilon: Added to variance for numerical stability when
        computing standard deviation.
    :returns: Tuple of two float arrays. The first is exponentially weighted
        moving means immediately after each event. The second is exponentially 
        weighted moving standard deviations immediately after each event.  
    '''

    momentum = 0.0
    mean = 0.0
    variance = 0.0

    means = np.empty_like(values)
    stds = np.empty_like(values)
    for i in range(values.shape[0]):
        retention = (1.0 - decay_rate) ** time_steps[i]

        new_momentum = momentum * retention + weights[i]

        means[i] = mean = \
            mean * (momentum / new_momentum) * retention + \
            values[i] * weights[i] / new_momentum

        deviation = values[i] - mean

        variance = \
            variance * (momentum / new_momentum) * retention + \
            np.square(deviation) * weights[i] / new_momentum

        momentum = new_momentum

        stds[i] = np.sqrt(variance + variance_epsilon)

    return means, stds

Sample usage:

means, stds = ewm(
    np.asarray([1., 2., 1., 2., 1., 2., 1., 2., 1., 2.]),
    np.asarray([1., 1., 2., 2., 1., 1., 2., 2., 1., 1.]),
    np.asarray([1., 1., 1., 1., 5., 5., 1., 1., 1., 1.]),
)
# [1.         1.50251256 1.33219236 1.5037909  1.2192523
#  1.5028668  1.4681632  1.50314792 1.47179905 1.50307306]
print(means)

# [0.001      0.35266091 0.34585958 0.39006561 0.30556194 
#  0.38630063 0.39249905 0.40020053 0.40503148 0.41104515]
print(stds)
\$\endgroup\$
0

1 Answer 1

1
\$\begingroup\$

You have multiple recursive (in the mathematical relation sense, not the computer science sense) expressions, notably on momentum, mean and variance. There is some minor vectorisation that can be done. I have tested this suggested code for correctness according to your provided examples. I've shown it to pull apart and illustrate what can be vectorised and what can't. I've also found accumulate to be a pain, because it seems to ignore the identity and dtype parameters.

Lo, there is hope: despite this looking as ugly as heck, in my testing it offers a ~180% speedup.

Suggested code

Docstrings omitted for brevity.

from timeit import timeit

import pandas as pd
import seaborn as sns
import numpy as np
from matplotlib import pyplot as plt
from numpy.random import default_rng


def ewm_old(values, time_steps, weights, decay_rate=0.01, variance_epsilon=1e-6):
    momentum = 0.0
    mean = 0.0
    variance = 0.0

    means = np.empty_like(values)
    stds = np.empty_like(values)
    for i in range(values.shape[0]):
        retention = (1.0 - decay_rate) ** time_steps[i]

        new_momentum = momentum * retention + weights[i]

        means[i] = mean = \
            mean * (momentum / new_momentum) * retention + \
            values[i] * weights[i] / new_momentum

        deviation = values[i] - mean

        variance = \
            variance * (momentum / new_momentum) * retention + \
            np.square(deviation) * weights[i] / new_momentum

        momentum = new_momentum

        stds[i] = np.sqrt(variance + variance_epsilon)

    return means, stds


def ewm_new(
    values: np.ndarray,
    time_steps: np.ndarray,
    weights: np.ndarray,
    decay_rate: float = 0.01,
    variance_epsilon: float = 1e-6,
) -> tuple[
    np.ndarray,
    np.ndarray,
]:
    def make_momentum(momentum: float, i: int) -> float:
        return momentum * retention[i-1] + weights[i-1]
    def make_means(mean: float, i: int) -> float:
        return mean*coefficients[i-1] + values[i-1]*offsets[i-1]
    def make_variances(variance: float, i: int) -> float:
        return variance*coefficients[i-1] + deviations[i-1]*offsets[i-1]

    retention = (1 - decay_rate)**time_steps
    momenta = np.frompyfunc(
        make_momentum, nin=2, nout=1, identity=0,
    ).accumulate(np.arange(len(values))).astype(np.float64)

    factors = momenta * retention / weights
    coefficients = factors / (factors + 1)
    offsets = 1 / (1 + factors)
    means = np.frompyfunc(
        make_means, nin=2, nout=1, identity=0,
    ).accumulate(np.arange(1+len(values)))[1:].astype(np.float64)

    deviations = (values - means)**2
    variances = np.frompyfunc(
        make_variances, nin=2, nout=1, identity=0,
    ).accumulate(np.arange(1+len(values)))[1:].astype(np.float64)

    stds = np.sqrt(variances + variance_epsilon)
    return means, stds


def test() -> None:
    rand = default_rng(seed=0)
    values, time_steps, weights = rand.random((3, 10))
    means, stds = ewm_new(values, time_steps, weights)

    assert np.allclose(
        means,
        np.array((
            0.63696169, 0.33792459, 0.09563421, 0.06077860, 0.28410923,
            0.38241663, 0.44761338, 0.51056811, 0.51504314, 0.56303342,
        )),
        rtol=0, atol=1e-6,
    )

    assert np.allclose(
        stds,
        np.array((
            0.00100000, 0.06149956, 0.05598960, 0.05115476, 0.29145246,
            0.34006516, 0.29895069, 0.28304818, 0.26340344, 0.27797243,
        )),
        rtol=0, atol=1e-6,
    )

    values     = np.array((1, 2, 1, 2, 1, 2, 1, 2, 1, 2), dtype=np.float64)
    time_steps = np.array((1, 1, 2, 2, 1, 1, 2, 2, 1, 1), dtype=np.float64)
    weights    = np.array((1, 1, 1, 1, 5, 5, 1, 1, 1, 1), dtype=np.float64)
    means, stds = ewm_new(values, time_steps, weights)

    assert np.allclose(
        means,
        (
            1.0000000,  1.50251256, 1.33219236, 1.50379090, 1.21925230,
            1.5028668,  1.46816320, 1.50314792, 1.47179905, 1.50307306,
        ),
        rtol=0, atol=1e-6,
    )

    assert np.allclose(
        stds,
        (
            0.00100000, 0.35266091, 0.34585958, 0.39006561, 0.30556194,
            0.38630063, 0.39249905, 0.40020053, 0.40503148, 0.41104515,
        ),
        rtol=0, atol=1e-6,
    )


def profile() -> None:
    rand = default_rng(seed=0)

    scale = np.round(10**np.linspace(0, 4, 150)).astype(int)
    times = []

    for n in scale:
        values, time_steps, weights = rand.random((3, n))

        for method in (ewm_old, ewm_new):
            t = timeit(lambda: method(values, time_steps, weights), number=1)
            times.append((method.__name__, n, t))

    df = pd.DataFrame(times, columns=('method', 'n', 'time'))
    sns.lineplot(data=df, x='n', y='time', hue='method')
    plt.show()


if __name__ == '__main__':
    test()
    profile()

profile

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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