# Exponentially-weighted moving mean and standard deviation of an irregularly-spaced weighted time series

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)


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()