Rename lagmat to sliding_window (#115)

Author: klane

tests/test_lagmat.py renamed to tests/test_sliding_window.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from torchts.utils.data import lagmat
+from torchts.utils.data import sliding_window
 
 
 @pytest.fixture
@@ -13,7 +13,7 @@ def tensor():
 @pytest.mark.parametrize("lag", [2, 5, [1, 2, 3], {1, 2, 3}, [1, 3, 5]])
 @pytest.mark.parametrize("horizon", [1, 2])
 def test_shape(tensor, lag, horizon):
-    x, y = lagmat(tensor, lag, horizon=horizon)
+    x, y = sliding_window(tensor, lag, horizon=horizon)
 
     if isinstance(lag, int):
         rows = len(tensor) - lag - horizon + 1
@@ -33,7 +33,7 @@ def test_shape(tensor, lag, horizon):
 @pytest.mark.parametrize("lag", [2, 5, [1, 2, 3], {1, 2, 3}, [1, 3, 5]])
 @pytest.mark.parametrize("horizon", [1, 2])
 def test_value(tensor, lag, horizon):
-    x, y = lagmat(tensor, lag, horizon=horizon)
+    x, y = sliding_window(tensor, lag, horizon=horizon)
 
     if isinstance(lag, int):
         for i in range(x.shape[0]):
@@ -49,10 +49,10 @@ def test_value(tensor, lag, horizon):
 @pytest.mark.parametrize("lag", ["1", 1.0, ["1"], [1, "2", 3], {1, 2.0, 3}])
 def test_non_int(tensor, lag):
     with pytest.raises(TypeError):
-        lagmat(tensor, lag)
+        sliding_window(tensor, lag)
 
 
 @pytest.mark.parametrize("lag", [-1, 0, [0, 1, 2], {0, 1, 2}, [-1, 1, 2]])
 def test_non_positive(tensor, lag):
     with pytest.raises(ValueError):
-        lagmat(tensor, lag)
+        sliding_window(tensor, lag)

torchts/utils/data.py

@@ -37,26 +37,6 @@ def concat(a, b):
     return torch.cat([a, b.unsqueeze(0)], dim=0)
 
 
-def lagmat(tensor, lags, horizon=1, dim=0, step=1):
-    is_int = isinstance(lags, int)
-    is_iter = isinstance(lags, Iterable) and all(isinstance(lag, int) for lag in lags)
-
-    if not is_int and not is_iter:
-        raise TypeError("lags must be of type int or Iterable[int]")
-
-    if (is_int and lags < 1) or (is_iter and any(lag < 1 for lag in lags)):
-        raise ValueError(f"lags must be positive but found {lags}")
-
-    if is_int:
-        data = tensor.unfold(dim, lags + horizon, step)
-        x, y = data[:, :lags], data[:, -1]
-    else:
-        data = tensor.unfold(dim, max(lags) + horizon, step)
-        x, y = data[:, [lag - 1 for lag in lags]], data[:, -1]
-
-    return x, y
-
-
 def load_dataset(dataset_dir, batch_size, val_batch_size=None, test_batch_size=None):
     if val_batch_size is None:
         val_batch_size = batch_size
@@ -108,3 +88,23 @@ def load_pickle(pickle_file):
         raise e
 
     return pickle_data
+
+
+def sliding_window(tensor, lags, horizon=1, dim=0, step=1):
+    is_int = isinstance(lags, int)
+    is_iter = isinstance(lags, Iterable) and all(isinstance(lag, int) for lag in lags)
+
+    if not is_int and not is_iter:
+        raise TypeError("lags must be of type int or Iterable[int]")
+
+    if (is_int and lags < 1) or (is_iter and any(lag < 1 for lag in lags)):
+        raise ValueError(f"lags must be positive but found {lags}")
+
+    if is_int:
+        data = tensor.unfold(dim, lags + horizon, step)
+        x, y = data[:, :lags], data[:, -1]
+    else:
+        data = tensor.unfold(dim, max(lags) + horizon, step)
+        x, y = data[:, [lag - 1 for lag in lags]], data[:, -1]
+
+    return x, y