"""Exponentially weighted moving average (EWMA) utilities and parameter conversion."""
import functools
import math
from dataclasses import field
import ezmsg.core as ez
import numpy as np
import numpy.typing as npt
import scipy.signal as sps
from array_api_compat import get_namespace, is_numpy_array
from ezmsg.baseproc import BaseStatefulTransformer, BaseTransformerUnit, processor_state, resolve_configured_stream_dim
from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
from ezmsg.util.messages.util import replace
from ezmsg.sigproc.util.array import np_finfo
from .util.deprecation import warn_axis_deprecated
def _ewma_mlx_metal_xp(data, axis_idx: int, zi, alpha: float, chunk_sizes: tuple[int, ...]):
"""Run EWMA through the MLX Metal helper while preserving scipy zi layout."""
import mlx.core as mx
from .util.ewma_mlx_metal import ewma_mlx_metal
zi = mx.asarray(zi, dtype=data.dtype)
last_data_axis = data.ndim - 1
last_zi_axis = zi.ndim - 1
x_mx = mx.moveaxis(data, axis_idx, last_data_axis) if axis_idx != last_data_axis else data
zi_mx = mx.moveaxis(zi, axis_idx, last_zi_axis) if axis_idx != last_zi_axis else zi
y_mx, zf_mx = ewma_mlx_metal(x_mx, alpha, zi_mx, chunk_sizes=chunk_sizes)
y = mx.moveaxis(y_mx, last_data_axis, axis_idx) if axis_idx != last_data_axis else y_mx
zf = mx.moveaxis(zf_mx, last_zi_axis, axis_idx) if axis_idx != last_zi_axis else zf_mx
return y, zf
def _tau_from_alpha(alpha: float, dt: float) -> float:
"""
Inverse of _alpha_from_tau. See that function for explanation.
"""
return -dt / np.log(1 - alpha)
def _alpha_from_tau(tau: float, dt: float) -> float:
"""
# https://en.wikipedia.org/wiki/Exponential_smoothing#Time_constant
:param tau: The amount of time for the smoothed response of a unit step function to reach
1 - 1/e approx-eq 63.2%.
:param dt: sampling period, or 1 / sampling_rate.
:return: alpha, the "fading factor" in exponential smoothing.
"""
return 1 - np.exp(-dt / tau)
[docs]
def ewma_step(sample: npt.NDArray, zi: npt.NDArray, alpha: float, beta: float | None = None):
"""
Do an exponentially weighted moving average step.
Args:
sample: The new sample.
zi: The output of the previous step.
alpha: Fading factor.
beta: Persisting factor. If None, it is calculated as 1-alpha.
Returns:
alpha * sample + beta * zi
"""
# Potential micro-optimization:
# Current: scalar-arr multiplication, scalar-arr multiplication, arr-arr addition
# Alternative: arr-arr subtraction, arr-arr multiplication, arr-arr addition
# return zi + alpha * (new_sample - zi)
beta = beta or (1 - alpha)
return alpha * sample + beta * zi
[docs]
class EWMA_Deprecated:
"""
Grabbed these methods from https://stackoverflow.com/a/70998068 and other answers in that topic,
but they ended up being slower than the scipy.signal.lfilter method.
Additionally, `compute` and `compute2` suffer from potential errors as the vector length increases
and beta**n approaches zero.
"""
[docs]
def __init__(self, alpha: float, max_len: int):
self.alpha = alpha
self.beta = 1 - alpha
self.prev: npt.NDArray | None = None
self.weights = np.empty((max_len + 1,), float)
self._precalc_weights(max_len)
self._step_func = functools.partial(ewma_step, alpha=self.alpha, beta=self.beta)
def _precalc_weights(self, n: int):
# (1-α)^0, (1-α)^1, (1-α)^2, ..., (1-α)^n
np.power(self.beta, np.arange(n + 1), out=self.weights)
[docs]
def compute(self, arr: npt.NDArray, out: npt.NDArray | None = None) -> npt.NDArray:
if out is None:
out = np.empty(arr.shape, arr.dtype)
n = arr.shape[0]
weights = self.weights[:n]
weights = np.expand_dims(weights, list(range(1, arr.ndim)))
# α*P0, α*P1, α*P2, ..., α*Pn
np.multiply(self.alpha, arr, out)
# α*P0/(1-α)^0, α*P1/(1-α)^1, α*P2/(1-α)^2, ..., α*Pn/(1-α)^n
np.divide(out, weights, out)
# α*P0/(1-α)^0, α*P0/(1-α)^0 + α*P1/(1-α)^1, ...
np.cumsum(out, axis=0, out=out)
# (α*P0/(1-α)^0)*(1-α)^0, (α*P0/(1-α)^0 + α*P1/(1-α)^1)*(1-α)^1, ...
np.multiply(out, weights, out)
# Add the previous output
if self.prev is None:
self.prev = arr[:1]
out += self.prev * np.expand_dims(self.weights[1 : n + 1], list(range(1, arr.ndim)))
self.prev = out[-1:]
return out
[docs]
def compute2(self, arr: npt.NDArray) -> npt.NDArray:
"""
Compute the Exponentially Weighted Moving Average (EWMA) of the input array.
Args:
arr: The input array to be smoothed.
Returns:
The smoothed array.
"""
n = arr.shape[0]
if n > len(self.weights):
self._precalc_weights(n)
weights = self.weights[:n][::-1]
weights = np.expand_dims(weights, list(range(1, arr.ndim)))
result = np.cumsum(self.alpha * weights * arr, axis=0)
result = result / weights
# Handle the first call when prev is unset
if self.prev is None:
self.prev = arr[:1]
result += self.prev * np.expand_dims(self.weights[1 : n + 1], list(range(1, arr.ndim)))
# Store the result back into prev
self.prev = result[-1]
return result
[docs]
def compute_sample(self, new_sample: npt.NDArray) -> npt.NDArray:
if self.prev is None:
self.prev = new_sample
self.prev = self._step_func(new_sample, self.prev)
return self.prev
[docs]
class EWMASettings(ez.Settings):
time_constant: float = 1.0
"""The amount of time for the smoothed response of a unit step function to reach 1 - 1/e approx-eq 63.2%."""
axis: str | None = None
""".. deprecated:: 3.8
Scheduled for removal in 4.0. The dimension messages accumulate along
now comes from :attr:`~ezmsg.util.messages.axisarray.AxisArray.stream_dim`;
see :mod:`ezmsg.sigproc.util.deprecation`."""
def __post_init__(self) -> None:
warn_axis_deprecated(self)
accumulate: bool = True
"""If True, update the EWMA state with each sample. If False, only apply
the current EWMA estimate without updating state (useful for inference
periods where you don't want to adapt statistics)."""
passthrough: bool = False
"""If True, return the input unchanged (identity) without touching the
EWMA. Unlike a very large time_constant -- which still applies a (stale)
baseline estimate -- passthrough leaves the data untouched. May be toggled
at runtime without resetting the filter state; see ``reset_on_resume`` for
what that means for the first message after the gap."""
reset_on_resume: bool = False
"""Whether switching ``passthrough`` back off discards the filter state.
The filter sees none of the samples that go by during passthrough, so ``zi``
describes an exponentially-weighted window that ended when passthrough was
switched on -- and the state cannot tell a 10 ms blip from a 10 minute
outage, since both resume identically. For a scaler that means z-scoring
post-gap data against pre-gap statistics.
False (the default) resumes from the preserved state, which is right for a
short blip and keeps an estimate that may have taken many ``time_constant``\\ s
to converge. True rebuilds from the first post-gap message instead: with the
bias correction below, that first output is exactly the first sample, and the
estimate re-converges over ``time_constant``. Prefer True where passthrough
may be left on long enough for the signal to drift, which is the case
:obj:`ezmsg.sigproc.binned_aggregate.BinnedAggregateTransformer` always
assumes.
Empty chunks are not gaps -- they carry no samples past the filter -- so they
never trigger this."""
mlx_metal_chunk_sizes: tuple[int, ...] = (32, 1024)
"""Allowable compile-time chunk sizes for EWMA Metal kernels. The smallest
size that fits the remaining samples is selected on each launch; otherwise
the largest size is repeated. Specializations compile lazily on first use.
Values must be in ``[1, 1024]``."""
def _bias_settle_count(alpha: float, dtype) -> float:
"""Cumulative sample count past which the bias correction rounds to exactly 1.
The correction divides by ``1 - (1-alpha)**t``, which rises monotonically
toward 1. Once ``(1-alpha)**t < eps/2`` the subtraction is a no-op in the
output dtype, so the division is by exactly 1.0 and can be skipped from then
on -- permanently, since ``t`` only grows. That is the steady state of any
long-running stream: it arrives after ~17 time constants regardless of rate.
A dtype NumPy cannot identify gets float64's epsilon, the most conservative
answer -- a too-small epsilon only forgoes an optimization, while a too-large
one would drop a correction that still mattered.
"""
finfo = np_finfo(dtype)
eps = float(finfo.eps if finfo is not None else np.finfo(np.float64).eps)
if not 0.0 < alpha < 1.0:
return math.inf
return math.log(0.5 * eps) / math.log1p(-alpha)
[docs]
@processor_state
class EWMAState:
alpha: float = field(default_factory=lambda: _alpha_from_tau(1.0, 1000.0))
zi: npt.NDArray | None = None
n_seen: int = 0
"""Cumulative sample count since reset, used to bias-correct the output."""
bias_settle_n: float = math.inf
"""``n_seen`` past which the bias correction is exactly 1; see :func:`_bias_settle_count`."""
bias_corr: npt.NDArray | None = None
"""Cached bias-correction divisor, in the message's namespace and dtype."""
bias_corr_for: tuple[int, int] | None = None
"""``(n_seen, n)`` that :attr:`bias_corr` was built for."""
[docs]
class EWMAUnit(BaseTransformerUnit[EWMASettings, AxisArray, AxisArray, EWMATransformer]):
SETTINGS = EWMASettings