add dim tests

Author: sarahmish

sigllm/primitives/prompting/timeseries_preprocessing.py

@@ -26,10 +26,14 @@ def rolling_window_sequences(X, window_size=500, step_size=100):
             * rolling window sequences.
             * first index value of each input sequence.
     """
+    if X.ndim == 1:
+        dim = 1
+    else:
+        dim = X.shape[1]
+
     index = range(len(X))
     out_X = list()
     X_index = list()
-    dim = X.shape[1]
 
     start = 0
     max_start = len(X) - window_size + 1

tests/primitives/prompting/test_timeseries_preprocessing.py

@@ -31,6 +31,35 @@ def test_rolling_window_sequences(values, window_size, step_size):
         np.array([0, 1, 2, 3, 4]),
         3,
         1,
+        1
+    )
+
+    result = rolling_window_sequences(values, window_size, step_size)
+
+    if len(result) != len(expected):
+        raise AssertionError('Tuples has different length')
+
+    for arr1, arr2 in zip(result, expected):
+        np.testing.assert_equal(arr1, arr2)
+
+
+def test_rolling_window_sequences_multivariate(window_size, step_size):
+    values = np.array([
+       [0.555, 2.345],
+       [1.501, 5.903],
+       [9.116, 3.068],
+       [7.432, 4.532]
+    ])
+
+    expected = (
+        np.array([
+            [[0.555, 2.345], [1.501, 5.903], [9.116, 3.068]],
+            [[1.501, 5.903], [9.116, 3.068], [7.432, 4.532]],
+        ]),
+        np.array([0, 1]),
+        3,
+        1,
+        2
     )
 
     result = rolling_window_sequences(values, window_size, step_size)