Add online op and new methods and class

pandas-dev · simonjayhawkins · Jun 12, 2021 · May 31, 2021 · May 31, 2021 · Jun 7, 2021
commit 9354bd0d261580c02459f8ed71f641d2c227c116
diff --git a/pandas/core/window/ewm.py b/pandas/core/window/ewm.py
@@ -711,7 +711,9 @@ def __init__(
             times=times,
             selection=selection,
         )
-        self._mean = EWMeanState(self._com, self.adjust, self.ignore_na)
+        self._mean = EWMeanState(
+            self._com, self.adjust, self.ignore_na, self.axis, obj.shape
+        )
         self.engine = engine
         self.engine_kwargs = engine_kwargs
 
@@ -747,7 +749,7 @@ def cov(
     def var(self, bias: bool = False, *args, **kwargs):
         return NotImplementedError
 
-    def mean(self, engine=None, engine_kwargs=None, update=None, update_deltas=None):
+    def mean(self, update=None, update_deltas=None):
         if update is not None:
             if self._mean.last_ewm is None:
                 raise ValueError(
@@ -756,11 +758,11 @@ def mean(self, engine=None, engine_kwargs=None, update=None, update_deltas=None)
             obj = np.concatenate(([self._mean.last_ewm], update.to_numpy()))
             result_from = 1
         else:
-            obj = self._selected_obj.to_numpy()
+            obj = self._selected_obj.astype(np.float64).to_numpy()
             result_from = 0
         if update_deltas is None:
             update_deltas = np.ones(max(len(obj) - 1, 0), dtype=np.float64)
-        ewma_func = generate_online_numba_ewma_func(engine_kwargs)
+        ewma_func = generate_online_numba_ewma_func(self.engine_kwargs)
         result = self._mean.run_ewm(obj, update_deltas, self.min_periods, ewma_func)
         result = self._selected_obj._constructor(result)
         return result.iloc[result_from:]
diff --git a/pandas/core/window/online.py b/pandas/core/window/online.py
@@ -40,56 +40,55 @@ def online_ewma(
         minimum_periods: int,
         old_wt_factor: float,
         new_wt: float,
-        old_wt: float,
+        old_wt: np.ndarray,
         adjust: bool,
         ignore_na: bool,
     ):
-        result = np.empty(len(values))
-
+        result = np.empty(values.shape)
         weighted_avg = values[0]
-        nobs = int(not np.isnan(weighted_avg))
-        result[0] = weighted_avg if nobs >= minimum_periods else np.nan
-
-        for j in range(1, len(values)):
-            cur = values[j]
-            is_observation = not np.isnan(cur)
-            nobs += is_observation
-            if not np.isnan(weighted_avg):
-
-                if is_observation or not ignore_na:
-
-                    # note that len(deltas) = len(vals) - 1 and deltas[i] is to be
-                    # used in conjunction with vals[i+1]
-                    old_wt *= old_wt_factor ** deltas[j - 1]
-                    if is_observation:
-
-                        # avoid numerical errors on constant series
-                        if weighted_avg != cur:
-                            weighted_avg = (
-                                (old_wt * weighted_avg) + (new_wt * cur)
-                            ) / (old_wt + new_wt)
-                        if adjust:
-                            old_wt += new_wt
-                        else:
-                            old_wt = 1.0
-            elif is_observation:
-                weighted_avg = cur
-
-            result[j] = weighted_avg if nobs >= minimum_periods else np.nan
+        nobs = (~np.isnan(weighted_avg)).astype(np.int64)
+        result[0] = np.where(nobs >= minimum_periods, weighted_avg, np.nan)
+
+        for i in range(1, len(values)):
+            cur = values[i]
+            is_observations = ~np.isnan(cur)
+            nobs += is_observations.astype(np.int64)
+            for j in range(len(cur)):
+                if not np.isnan(weighted_avg[j]):
+                    if is_observations[j] or not ignore_na:
+                        # note that len(deltas) = len(vals) - 1 and deltas[i] is to be
+                        # used in conjunction with vals[i+1]
+                        old_wt[j] *= old_wt_factor ** deltas[j - 1]
+                        if is_observations[j]:
+                            # avoid numerical errors on constant series
+                            if weighted_avg[j] != cur[j]:
+                                weighted_avg[j] = (
+                                    (old_wt[j] * weighted_avg[j]) + (new_wt * cur[j])
+                                ) / (old_wt[j] + new_wt)
+                            if adjust:
+                                old_wt[j] += new_wt
+                            else:
+                                old_wt[j] = 1.0
+                elif is_observations[j]:
+                    weighted_avg[j] = cur[j]
+
+            result[i] = np.where(nobs >= minimum_periods, weighted_avg, np.nan)
 
         return result, old_wt
 
     return online_ewma
 
 
 class EWMeanState:
-    def __init__(self, com, adjust, ignore_na):
+    def __init__(self, com, adjust, ignore_na, axis, shape):
         alpha = 1.0 / (1.0 + com)
-        self.old_wt_factor = 1.0 - alpha
-        self.new_wt = 1.0 if adjust else alpha
-        self.old_wt = 1.0
+        self.axis = axis
+        self.shape = shape
         self.adjust = adjust
         self.ignore_na = ignore_na
+        self.new_wt = 1.0 if adjust else alpha
+        self.old_wt_factor = 1.0 - alpha
+        self.old_wt = np.ones(self.shape[self.axis - 1])
         self.last_ewm = None
 
     def run_ewm(self, weighted_avg, deltas, min_periods, ewm_func):
@@ -108,5 +107,5 @@ def run_ewm(self, weighted_avg, deltas, min_periods, ewm_func):
         return result
 
     def reset(self):
-        self.old_wt = 1
+        self.old_wt = np.ones(self.shape[self.axis - 1])
         self.last_ewm = None