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():