Closed
Description
tis is an implementation of online ewma. its self contained, copying mostly the pandas ewma implementation (and exposing the outputs).
code
import typing
import numba
import numpy as np
from numpy import nan
def ewma(
x: np.ndarray,
*,
alpha: float,
min_periods: int,
ignore_na: bool,
) -> typing.Tuple[
np.ndarray,
np.ndarray,
np.ndarray
]:
assert isinstance(x, np.ndarray)
assert len(x.shape) == 1
assert x.dtype == np.float64
n: np.ndarray = np.empty_like(x, dtype=np.int_)
w: np.ndarray = np.empty_like(x)
y: np.ndarray = np.empty_like(x)
if x.shape[0]:
y_0: float = x[0]
is_observation: bool = y_0 == y_0
n_0: int = int(is_observation)
w_0: float = 1.0
y[0] = y_0 if n_0 >= min_periods else nan
n[0] = n_0
w[0] = w_0
_inc_ewma(
x,
n,
w,
y,
alpha,
min_periods,
ignore_na,
n_0,
w_0,
y_0,
1
)
return n, w, y
def inc_ewma(
x: np.ndarray,
*,
alpha: float,
min_periods: int,
ignore_na: bool,
n_minus_1: int,
w_minus_1: float,
y_minus_1: float
) -> typing.Tuple[
np.ndarray,
np.ndarray,
np.ndarray
]:
assert isinstance(x, np.ndarray)
assert len(x.shape) == 1
assert x.dtype == np.float64
n: np.ndarray = np.empty_like(x, dtype=np.int_)
w: np.ndarray = np.empty_like(x)
y: np.ndarray = np.empty_like(x)
_inc_ewma(
x,
n,
w,
y,
alpha,
min_periods,
ignore_na,
n_minus_1,
w_minus_1,
y_minus_1,
0
)
return n, w, y
@numba.njit # type: ignore
def _inc_ewma(
x: np.ndarray,
n: np.ndarray,
w: np.ndarray,
y: np.ndarray,
alpha: float,
min_periods: int,
ignore_na: bool,
n_i: int,
w_i: float,
y_i: float,
i: int
) -> None:
beta: float = 1.0 - alpha
for i in range(i, len(x)):
x_i: float = x[i]
is_observation = x_i == x_i
n_i += is_observation
if y_i == y_i:
if is_observation or not ignore_na:
w_i *= beta
if is_observation:
# avoid numerical errors on constant series
if y_i != x_i:
y_i = ((w_i * y_i) + x_i) / (w_i + 1.0)
w_i += 1.0
elif is_observation:
y_i = x_i
y[i] = y_i if n_i >= min_periods else nan
n[i] = n_i
w[i] = w_i
tests
import unittest
import numpy as np
import pandas as pd
from incewma import ewma, inc_ewma
class IncEwmaTestCase(unittest.TestCase):
def test_ewma(self) -> None:
alpha = 0.5
min_periods = 3
ignore_na = False
x = pd.Series(np.arange(10, dtype=np.float64) + 1.7)
y: np.ndarray = x.ewm(
alpha=alpha,
min_periods=min_periods,
adjust=True,
ignore_na=ignore_na
).mean().to_numpy()
n_all, w_all, y_all = ewma(
x.to_numpy(),
alpha=alpha,
min_periods=min_periods,
ignore_na=ignore_na
)
np.testing.assert_equal(np.arange(1, x.shape[0] + 1), n_all)
np.testing.assert_allclose(y, y_all)
k = 3
n_k, w_k, y_k = inc_ewma(
x.to_numpy()[-k:],
alpha=alpha,
min_periods=min_periods,
ignore_na=ignore_na,
n_minus_1=n_all[-(k + 1)],
w_minus_1=w_all[-(k + 1)],
y_minus_1=y_all[-(k + 1)]
)
np.testing.assert_equal(n_all[-k:], n_k)
np.testing.assert_allclose(w_all[-k:], w_k)
np.testing.assert_allclose(y[-k:], y_k)
def test_ewma_empty(self) -> None:
alpha = 0.5
min_periods = 3
ignore_na = False
x = pd.Series(np.empty((0,), dtype=np.float64))
y: np.ndarray = x.ewm(
alpha=alpha,
min_periods=min_periods,
adjust=True,
ignore_na=ignore_na
).mean().to_numpy()
np.testing.assert_allclose(np.empty((0,), dtype=np.float64), y)
n_all, w_all, y_all = ewma(
x.to_numpy(),
alpha=alpha,
min_periods=min_periods,
ignore_na=ignore_na
)
np.testing.assert_equal(np.empty((0,), dtype=np.int_), n_all)
np.testing.assert_allclose(np.empty((0,), dtype=np.float64), w_all)
np.testing.assert_allclose(np.empty((0,), dtype=np.float64), y_all)
n_k, w_k, y_k = inc_ewma(
x.to_numpy(),
alpha=alpha,
min_periods=min_periods,
ignore_na=ignore_na,
n_minus_1=1,
w_minus_1=1.0,
y_minus_1=1.0
)
np.testing.assert_equal(np.empty((0,), dtype=np.int_), n_k)
np.testing.assert_allclose(np.empty((0,), dtype=np.float64), w_k)
np.testing.assert_allclose(np.empty((0,), dtype=np.float64), y_k)
Metadata
Metadata
Assignees
Labels
No labels