From e341f9cb548bc2f2fc8f0d64430a76b7c6ecac6c Mon Sep 17 00:00:00 2001
From: Kashif Rasul <kashif.rasul@gmail.com>
Date: Sun, 17 Nov 2019 22:46:57 +0100
Subject: [PATCH] initial test for get_lagged_subseq

---
 test/model/deepar/test_lags.py | 52 ++++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)
 create mode 100644 test/model/deepar/test_lags.py

diff --git a/test/model/deepar/test_lags.py b/test/model/deepar/test_lags.py
new file mode 100644
index 0000000..5e6c70e
--- /dev/null
+++ b/test/model/deepar/test_lags.py
@@ -0,0 +1,52 @@
+# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the "license" file accompanying this file. This file is distributed
+# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+# express or implied. See the License for the specific language governing
+# permissions and limitations under the License.
+
+# Standard library imports
+import itertools
+
+# Third-party imports
+import torch
+
+# First-party imports
+from pts.model.deepar._network import DeepARTrainingNetwork
+
+
+def test_lagged_subsequences():
+    N = 8
+    T = 96
+    C = 2
+    lags = [1, 2, 3, 24, 48]
+    I = len(lags)
+    sequence = torch.randn((N, T, C))
+    S = 48
+
+    # (batch_size, sub_seq_len, target_dim, num_lags)
+    lagged_subsequences = DeepARTrainingNetwork.get_lagged_subsequences(
+        sequence=sequence,
+        sequence_length=sequence.shape[1],
+        indices=lags,
+        subsequences_length=S,
+    )
+
+    assert (N, S, C, I) == lagged_subsequences.shape
+
+    # checks that lags value behave as described as in the get_lagged_subsequences contract
+    for i, j, k in itertools.product(range(N), range(S), range(I)):
+        assert (
+            (
+                lagged_subsequences[i, j, :, k]
+                == sequence[i, -lags[k] - S + j, :]
+            )
+            .numpy()
+            .all()
+        )