add stack method (#24)

Author: jmoralez

coreforecast/lag_transforms.py

@@ -20,14 +20,16 @@
 
 import abc
 import copy
-from typing import Callable, Optional
+from typing import Callable, Optional, Sequence
 
 import numpy as np
 
 from .grouped_array import GroupedArray
 
 
 class _BaseLagTransform(abc.ABC):
+    stats_: np.ndarray
+
     @abc.abstractmethod
     def transform(self, ga: GroupedArray) -> np.ndarray:
         """Apply the transformation by group.
@@ -53,6 +55,20 @@ def update(self, ga: GroupedArray) -> np.ndarray:
     def take(self, _idxs: np.ndarray) -> "_BaseLagTransform":
         return self
 
+    @staticmethod
+    def stack(transforms: Sequence["_BaseLagTransform"]) -> "_BaseLagTransform":
+        first_tfm = transforms[0]
+        if not hasattr(first_tfm, "stats_"):
+            # transform doesn't save state, we can return any of them
+            return first_tfm
+        out = copy.deepcopy(first_tfm)
+        if first_tfm.stats_.ndim == 1:
+            concat_fn = np.hstack
+        else:
+            concat_fn = np.vstack
+        out.stats_ = concat_fn([tfm.stats_ for tfm in transforms])
+        return out
+
 
 class Lag(_BaseLagTransform):
     """Simple lag operator
@@ -396,15 +412,15 @@ def __init__(self, lag: int, alpha: float):
 
     def transform(self, ga: GroupedArray) -> np.ndarray:
         out = ga._exponentially_weighted_transform("Mean", self.lag, self.alpha)
-        self.ewm_ = out[ga.indptr[1:] - 1]
+        self.stats_ = out[ga.indptr[1:] - 1]
         return out
 
     def update(self, ga: GroupedArray) -> np.ndarray:
         x = ga._index_from_end(self.lag - 1)
-        self.ewm_ = self.alpha * x + (1 - self.alpha) * self.ewm_
-        return self.ewm_
+        self.stats_ = self.alpha * x + (1 - self.alpha) * self.stats_
+        return self.stats_
 
     def take(self, idxs: np.ndarray) -> "ExponentiallyWeightedMean":
         out = copy.deepcopy(self)
-        out.ewm_ = out.ewm_[idxs].copy()
+        out.stats_ = out.stats_[idxs].copy()
         return out

coreforecast/scalers.py

@@ -1,6 +1,6 @@
 import copy
 import ctypes
-from typing import List, Optional
+from typing import List, Optional, Sequence
 
 import numpy as np
 
@@ -205,6 +205,12 @@ def take(self, idxs: np.ndarray) -> "_BaseLocalScaler":
         out.stats_ = self.stats_[idxs].copy()
         return out
 
+    @staticmethod
+    def stack(scalers: Sequence["_BaseLocalScaler"]) -> "_BaseLocalScaler":
+        out = copy.deepcopy(scalers[0])
+        out.stats_ = np.vstack([sc.stats_ for sc in scalers])
+        return out
+
 
 class LocalMinMaxScaler(_BaseLocalScaler):
     """Scale each group to the [0, 1] interval"""
@@ -354,6 +360,12 @@ def take(self, idxs: np.ndarray) -> "Difference":
         out.tails_ = self.tails_[idxs].copy()
         return out
 
+    @staticmethod
+    def stack(scalers: Sequence["Difference"]) -> "Difference":
+        out = Difference(scalers[0].d)
+        out.tails_ = np.hstack([sc.tails_ for sc in scalers])
+        return out
+
 
 class AutoDifferences:
     """Find and apply the optimal number of differences to each group.
@@ -447,6 +459,16 @@ def take(self, idxs: np.ndarray) -> "AutoDifferences":
         out.tails_ = [tail[idxs].copy() for tail in self.tails_]
         return out
 
+    @staticmethod
+    def stack(scalers: Sequence["AutoDifferences"]) -> "AutoDifferences":
+        out = copy.deepcopy(scalers[0])
+        out.diffs_ = np.hstack([sc.diffs_ for sc in scalers])
+        out.tails_ = [
+            np.hstack([sc.tails_[i] for sc in scalers])
+            for i in range(len(scalers[0].tails_))
+        ]
+        return out
+
 
 class AutoSeasonalDifferences(AutoDifferences):
     """Find and apply the optimal number of seasonal differences to each group.
@@ -604,3 +626,14 @@ def take(self, idxs: np.ndarray) -> "AutoSeasonalityAndDifferences":
         out.diffs_ = [diffs[idxs].copy() for diffs in self.diffs_]
         out.tails_ = [tail[idxs].copy() for tail in self.tails_]
         return out
+
+    @staticmethod
+    def stack(
+        scalers: Sequence["AutoSeasonalityAndDifferences"],
+    ) -> "AutoSeasonalityAndDifferences":
+        out = AutoSeasonalityAndDifferences(
+            scalers[0].max_season_length, scalers[0].max_diffs, scalers[0].n_seasons
+        )
+        out.diffs_ = [np.hstack([sc.diffs_[i] for sc in scalers]) for i in range(out.max_diffs)]
+        out.tails_ = [np.hstack([sc.tails_[i] for sc in scalers]) for i in range(out.max_diffs)]
+        return out

tests/test_lag_transforms.py

@@ -116,6 +116,12 @@ def test_correctness(data, comb, dtype):
     wres = transform(data, indptr, True, lag - 1, wf, *args)
     cres = cobj.update(ga)
     np.testing.assert_allclose(wres, cres, atol=atol, rtol=rtol)
+    # stack
+    combined = cobj.stack([cobj, cobj])
+    if hasattr(cobj, "stats_"):
+        assert combined.stats_.shape[0] == 2 * cobj.stats_.shape[0]
+    else:
+        assert combined is cobj
 
 
 @pytest.mark.parametrize("window_type", ["rolling", "seasonal_rolling", "expanding"])