From bae53cfcfa6d3813fa15472c8d4cfa29dd817a10 Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Mon, 18 Nov 2019 16:05:31 +0100
Subject: [PATCH] added get_lagged_subsequences

---
 pts/model/deepar/deepar_network.py | 22 +++++++++++++++++++++-
 test/model/deepar/test_lags.py     |  2 +-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/pts/model/deepar/deepar_network.py b/pts/model/deepar/deepar_network.py
index 1742bd0..05b33f6 100644
--- a/pts/model/deepar/deepar_network.py
+++ b/pts/model/deepar/deepar_network.py
@@ -1,6 +1,10 @@
+from typing import List
+
 import torch
 import torch.nn as nn
 
+import numpy as np
+
 from pts.modules import DistributionOutput, MeanScaler, NOPScaler, FeatureEmbedder
 
 
@@ -70,7 +74,23 @@ class DeepARNetwork(nn.Module):
         sequence_length: int,
         indices: List[int],
         subsequences_length: int = 1,
-    ) -> torch.Tensor
+    ) -> torch.Tensor:
+
+        assert max(indices) + subsequences_length <= sequence_length, (
+            f"lags cannot go further than history length, found lag {max(indices)} "
+            f"while history length is only {sequence_length}"
+        )
+        assert all(lag_index >= 0 for lag_index in indices)
+
+        lagged_values = []
+        for lag_index in indices:
+            begin_index = -lag_index - subsequences_length
+            end_index = -lag_index if lag_index > 0 else None
+            lagged_values.append(
+                sequence[:,begin_index:end_index,...]
+            )
+        return torch.stack(lagged_values, dim=-1)
+
 
 class DeepARTrainingNetwork(DeepARNetwork):
     pass
\ No newline at end of file
diff --git a/test/model/deepar/test_lags.py b/test/model/deepar/test_lags.py
index 5e6c70e..17d6362 100644
--- a/test/model/deepar/test_lags.py
+++ b/test/model/deepar/test_lags.py
@@ -18,7 +18,7 @@ import itertools
 import torch
 
 # First-party imports
-from pts.model.deepar._network import DeepARTrainingNetwork
+from pts.model.deepar import DeepARTrainingNetwork
 
 
 def test_lagged_subsequences():