From 4e97483fab8415b3dbcee7ca270ea19350e0a734 Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Tue, 30 Mar 2021 16:55:13 +0200 Subject: [PATCH] WIP: causal model (#31) * intial causal model * added training * initial inference * formatting * added readme to model * key is called control * added schema of the model * formatting * get rid of confusing figure. * typi * more typos * clarification * added some examples * fix covariate convention --- pts/model/causal_deepar/README.md | 53 ++ pts/model/causal_deepar/__init__.py | 2 + .../causal_deepar/causal_deepar_estimator.py | 280 +++++++++ .../causal_deepar/causal_deepar_network.py | 570 ++++++++++++++++++ 4 files changed, 905 insertions(+) create mode 100644 pts/model/causal_deepar/README.md create mode 100644 pts/model/causal_deepar/__init__.py create mode 100644 pts/model/causal_deepar/causal_deepar_estimator.py create mode 100644 pts/model/causal_deepar/causal_deepar_network.py diff --git a/pts/model/causal_deepar/README.md b/pts/model/causal_deepar/README.md new file mode 100644 index 0000000..377b9a0 --- /dev/null +++ b/pts/model/causal_deepar/README.md @@ -0,0 +1,53 @@ +# Causal `DeepAR` + +Causal `DeepAR` model augments the `DeepAR` model by incorporating a causal structure via a `control` time-dependent signal. + +The main assumption of this model is that the `target` at time `t` depends not only on the covariates up till time `t` and but also on `control` till time `t`. Thus we encode this structure via a model which now adds a `control_output` distribution layer. The main assumption we have is that at training time our dataset of time series now have an additional `control` key with the corresponding 1-d control-variate values (with the array being as large as `target`). + +We translate this structure into the following schematics of Causal-`DeepAR` at time `t`: + +``` +┌─────────┐ ┌────────┐ +│control_t│──────┐ │target_t│ +└─────────┘ │ └────────┘ + ▲ │ ▲ + │ │ │ + log-prob/ │ log-prob/ + sample │ sample + │ ground-truth/ │ +┌────────┐ sample/ ┌───────┐ +│control │ do │target │ +│ dist │ └──────▶│ dist │ +└────────┘ └───────┘ + ▲ ▲ + └────────────┬──────────┘ + │ + ┌────┐ + ───h_t-1─▶│RNN │───h_t──▶ + └────┘ + ▲ + ┌─────┴────┐ + │target_t-1│ + └─────┬────┘ + ┌─────┴─────┐ + │control_t-1│ + └─────┬─────┘ + ┌───┴──┐ + │cov_t │ + └──────┘ +``` + + +The model is trained as per the DeepAR assumption, which mainly implies that the covariates and `control` are known for all the time points while training and only the covariates are known for all the time points that we wish to forecast for. + +In terms of the `control` at inference time, we now have two choices: + +1. We can predict by setting the `control` array to `np.Nan` for the duration of the prediction length in which case the model will sample values from the `control_output.distribution` and feed it auto-regressively to the `target` head and RNN. + +2. On the other hand we can in the prediction window "do" an intervention by setting the values of `control` to some fixed value of our choosing for the time steps in the future we are interested in. In this case this model will "break" the causal connections and just use the supplied values and feed those to the `target` head and RNN at the appropriate time points. + +## Possible uses + +* `target` can be sales and `control` can be discounts +* `target` can be sales and `control` can be the temperature + diff --git a/pts/model/causal_deepar/__init__.py b/pts/model/causal_deepar/__init__.py new file mode 100644 index 0000000..86e1d4d --- /dev/null +++ b/pts/model/causal_deepar/__init__.py @@ -0,0 +1,2 @@ +from .causal_deepar_estimator import CausalDeepAREstimator +from .causal_deepar_network import CausalDeepARNetwork, CausalDeepARTrainingNetwork diff --git a/pts/model/causal_deepar/causal_deepar_estimator.py b/pts/model/causal_deepar/causal_deepar_estimator.py new file mode 100644 index 0000000..55cc04a --- /dev/null +++ b/pts/model/causal_deepar/causal_deepar_estimator.py @@ -0,0 +1,280 @@ +from typing import List, Optional + +import numpy as np +import torch +import torch.nn as nn + +from gluonts.core.component import validated +from gluonts.dataset.field_names import FieldName +from gluonts.time_feature import ( + TimeFeature, + get_lags_for_frequency, + time_features_from_frequency_str, +) +from gluonts.transform import ( + Transformation, + Chain, + RemoveFields, + SetField, + AsNumpyArray, + AddObservedValuesIndicator, + AddTimeFeatures, + AddAgeFeature, + VstackFeatures, + InstanceSplitter, + ValidationSplitSampler, + TestSplitSampler, + ExpectedNumInstanceSampler, +) +from gluonts.torch.support.util import copy_parameters +from gluonts.torch.model.predictor import PyTorchPredictor +from gluonts.torch.modules.distribution_output import DistributionOutput +from gluonts.model.predictor import Predictor + +from pts.model.utils import get_module_forward_input_names +from pts import Trainer +from pts.model import PyTorchEstimator +from pts.modules import StudentTOutput + +from .causal_deepar_network import ( + CausalDeepARTrainingNetwork, + CausalDeepARPredictionNetwork, +) + + +class CausalDeepAREstimator(PyTorchEstimator): + @validated() + def __init__( + self, + freq: str, + prediction_length: int, + input_size: int, + trainer: Trainer = Trainer(), + context_length: Optional[int] = None, + num_layers: int = 2, + num_cells: int = 40, + cell_type: str = "LSTM", + dropout_rate: float = 0.1, + use_feat_dynamic_real: bool = False, + use_feat_dynamic_cat: bool = False, + use_feat_static_cat: bool = False, + use_feat_static_real: bool = False, + cardinality: Optional[List[int]] = None, + embedding_dimension: Optional[List[int]] = None, + distr_output: DistributionOutput = StudentTOutput(), + control_output: DistributionOutput = StudentTOutput(), + scaling: bool = True, + lags_seq: Optional[List[int]] = None, + time_features: Optional[List[TimeFeature]] = None, + num_parallel_samples: int = 100, + dtype: np.dtype = np.float32, + ) -> None: + super().__init__(trainer=trainer) + + self.freq = freq + self.context_length = ( + context_length if context_length is not None else prediction_length + ) + self.prediction_length = prediction_length + self.distr_output = distr_output + self.control_output = control_output + self.distr_output.dtype = dtype + self.input_size = input_size + self.num_layers = num_layers + self.num_cells = num_cells + self.cell_type = cell_type + self.dropout_rate = dropout_rate + self.use_feat_dynamic_real = use_feat_dynamic_real + self.use_feat_dynamic_cat = use_feat_dynamic_cat + self.use_feat_static_cat = use_feat_static_cat + self.use_feat_static_real = use_feat_static_real + self.cardinality = cardinality if cardinality and use_feat_static_cat else [1] + self.embedding_dimension = ( + embedding_dimension + if embedding_dimension is not None + else [min(50, (cat + 1) // 2) for cat in self.cardinality] + ) + self.scaling = scaling + self.lags_seq = ( + lags_seq + if lags_seq is not None + else get_lags_for_frequency(freq_str=freq, lag_ub=self.context_length) + ) + self.time_features = ( + time_features + if time_features is not None + else time_features_from_frequency_str(self.freq) + ) + + self.history_length = self.context_length + max(self.lags_seq) + + self.num_parallel_samples = num_parallel_samples + + self.train_sampler = ExpectedNumInstanceSampler( + num_instances=1.0, min_future=prediction_length + ) + self.validation_sampler = ValidationSplitSampler(min_future=prediction_length) + + def create_transformation(self) -> Transformation: + remove_field_names = [] + if not self.use_feat_static_real: + remove_field_names.append(FieldName.FEAT_STATIC_REAL) + if not self.use_feat_dynamic_real: + remove_field_names.append(FieldName.FEAT_DYNAMIC_REAL) + if not self.use_feat_dynamic_cat: + remove_field_names.append(FieldName.FEAT_DYNAMIC_CAT) + + return Chain( + [RemoveFields(field_names=remove_field_names)] + + ( + [SetField(output_field=FieldName.FEAT_STATIC_CAT, value=[0])] + if not self.use_feat_static_cat + else [] + ) + + ( + [SetField(output_field=FieldName.FEAT_STATIC_REAL, value=[0.0])] + if not self.use_feat_static_real + else [] + ) + + [ + AsNumpyArray( + field=FieldName.FEAT_STATIC_CAT, + expected_ndim=1, + dtype=np.long, + ), + AsNumpyArray( + field=FieldName.FEAT_STATIC_REAL, + expected_ndim=1, + dtype=self.dtype, + ), + AsNumpyArray( + field=FieldName.TARGET, + # in the following line, we add 1 for the time dimension + expected_ndim=1 + len(self.distr_output.event_shape), + dtype=self.dtype, + ), + AsNumpyArray( + field="control", + # in the following line, we add 1 for the time dimension + expected_ndim=1 + len(self.control_output.event_shape), + dtype=self.dtype, + ), + AddObservedValuesIndicator( + target_field=FieldName.TARGET, + output_field=FieldName.OBSERVED_VALUES, + dtype=self.dtype, + ), + AddTimeFeatures( + start_field=FieldName.START, + target_field=FieldName.TARGET, + output_field=FieldName.FEAT_TIME, + time_features=self.time_features, + pred_length=self.prediction_length, + ), + AddAgeFeature( + target_field=FieldName.TARGET, + output_field=FieldName.FEAT_AGE, + pred_length=self.prediction_length, + log_scale=True, + dtype=self.dtype, + ), + VstackFeatures( + output_field=FieldName.FEAT_TIME, + input_fields=[FieldName.FEAT_TIME, FieldName.FEAT_AGE] + + ( + [FieldName.FEAT_DYNAMIC_REAL] + if self.use_feat_dynamic_real + else [] + ) + + ( + [FieldName.FEAT_DYNAMIC_CAT] + if self.use_feat_dynamic_cat + else [] + ), + ), + ] + ) + + def create_instance_splitter(self, mode: str): + assert mode in ["training", "validation", "test"] + + instance_sampler = { + "training": self.train_sampler, + "validation": self.validation_sampler, + "test": TestSplitSampler(), + }[mode] + + return InstanceSplitter( + target_field=FieldName.TARGET, + is_pad_field=FieldName.IS_PAD, + start_field=FieldName.START, + forecast_start_field=FieldName.FORECAST_START, + instance_sampler=instance_sampler, + past_length=self.history_length, + future_length=self.prediction_length, + time_series_fields=[ + FieldName.FEAT_TIME, + FieldName.OBSERVED_VALUES, + "control", + ], + ) + + def create_training_network( + self, device: torch.device + ) -> CausalDeepARTrainingNetwork: + return CausalDeepARTrainingNetwork( + input_size=self.input_size, + num_layers=self.num_layers, + num_cells=self.num_cells, + cell_type=self.cell_type, + history_length=self.history_length, + context_length=self.context_length, + prediction_length=self.prediction_length, + distr_output=self.distr_output, + control_output=self.control_output, + dropout_rate=self.dropout_rate, + cardinality=self.cardinality, + embedding_dimension=self.embedding_dimension, + lags_seq=self.lags_seq, + scaling=self.scaling, + dtype=self.dtype, + ).to(device) + + def create_predictor( + self, + transformation: Transformation, + trained_network: nn.Module, + device: torch.device, + ) -> Predictor: + prediction_network = CausalDeepARPredictionNetwork( + num_parallel_samples=self.num_parallel_samples, + input_size=self.input_size, + num_layers=self.num_layers, + num_cells=self.num_cells, + cell_type=self.cell_type, + history_length=self.history_length, + context_length=self.context_length, + prediction_length=self.prediction_length, + distr_output=self.distr_output, + control_output=self.control_output, + dropout_rate=self.dropout_rate, + cardinality=self.cardinality, + embedding_dimension=self.embedding_dimension, + lags_seq=self.lags_seq, + scaling=self.scaling, + dtype=self.dtype, + ).to(device) + + copy_parameters(trained_network, prediction_network) + input_names = get_module_forward_input_names(prediction_network) + prediction_splitter = self.create_instance_splitter("test") + + return PyTorchPredictor( + input_transform=transformation + prediction_splitter, + input_names=input_names, + prediction_net=prediction_network, + batch_size=self.trainer.batch_size, + freq=self.freq, + prediction_length=self.prediction_length, + device=device, + ) diff --git a/pts/model/causal_deepar/causal_deepar_network.py b/pts/model/causal_deepar/causal_deepar_network.py new file mode 100644 index 0000000..64a9c2c --- /dev/null +++ b/pts/model/causal_deepar/causal_deepar_network.py @@ -0,0 +1,570 @@ +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn as nn +from torch.distributions import Distribution + +from gluonts.core.component import validated +from gluonts.torch.modules.distribution_output import DistributionOutput +from pts.model import weighted_average +from pts.modules import MeanScaler, NOPScaler, FeatureEmbedder + + +def prod(xs): + p = 1 + for x in xs: + p *= x + return p + + +class CausalDeepARNetwork(nn.Module): + @validated() + def __init__( + self, + input_size: int, + num_layers: int, + num_cells: int, + cell_type: str, + history_length: int, + context_length: int, + prediction_length: int, + distr_output: DistributionOutput, + control_output: DistributionOutput, + dropout_rate: float, + cardinality: List[int], + embedding_dimension: List[int], + lags_seq: List[int], + scaling: bool = True, + dtype: np.dtype = np.float32, + ) -> None: + super().__init__() + self.num_layers = num_layers + self.num_cells = num_cells + self.cell_type = cell_type + self.history_length = history_length + self.context_length = context_length + self.prediction_length = prediction_length + self.dropout_rate = dropout_rate + self.cardinality = cardinality + self.embedding_dimension = embedding_dimension + self.num_cat = len(cardinality) + self.scaling = scaling + self.dtype = dtype + + self.lags_seq = lags_seq + + self.distr_output = distr_output + self.control_output = control_output + + rnn = {"LSTM": nn.LSTM, "GRU": nn.GRU}[self.cell_type] + self.rnn = rnn( + input_size=input_size, + hidden_size=num_cells, + num_layers=num_layers, + dropout=dropout_rate, + batch_first=True, + ) + + self.target_shape = distr_output.event_shape + + self.proj_distr_args = distr_output.get_args_proj(num_cells + 1) + self.proj_control_args = control_output.get_args_proj(num_cells) + + self.embedder = FeatureEmbedder( + cardinalities=cardinality, embedding_dims=embedding_dimension + ) + + if scaling: + self.scaler = MeanScaler(keepdim=True) + self.control_scaler = NOPScaler(keepdim=True) + else: + self.scaler = NOPScaler(keepdim=True) + self.control_scaler = NOPScaler(keepdim=True) + + @staticmethod + def get_lagged_subsequences( + sequence: torch.Tensor, + sequence_length: int, + indices: List[int], + subsequences_length: int = 1, + ) -> torch.Tensor: + """ + Returns lagged subsequences of a given sequence. + Parameters + ---------- + sequence : Tensor + the sequence from which lagged subsequences should be extracted. + Shape: (N, T, C). + sequence_length : int + length of sequence in the T (time) dimension (axis = 1). + indices : List[int] + list of lag indices to be used. + subsequences_length : int + length of the subsequences to be extracted. + Returns + -------- + lagged : Tensor + a tensor of shape (N, S, C, I), where S = subsequences_length and + I = len(indices), containing lagged subsequences. Specifically, + lagged[i, j, :, k] = sequence[i, -indices[k]-S+j, :]. + """ + 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) + + def unroll_encoder( + self, + feat_static_cat: torch.Tensor, # (batch_size, num_features) + feat_static_real: torch.Tensor, # (batch_size, num_features) + past_time_feat: torch.Tensor, # (batch_size, history_length, num_features) + past_target: torch.Tensor, # (batch_size, history_length, *target_shape) + past_observed_values: torch.Tensor, # (batch_size, history_length, *target_shape) + past_control: torch.Tensor, + future_control: Optional[torch.Tensor] = None, + future_time_feat: Optional[ + torch.Tensor + ] = None, # (batch_size, prediction_length, num_features) + future_target: Optional[ + torch.Tensor + ] = None, # (batch_size, prediction_length, *target_shape) + ) -> Tuple[torch.Tensor, Union[torch.Tensor, List], torch.Tensor, torch.Tensor]: + + if future_time_feat is None or future_target is None or future_control is None: + time_feat = past_time_feat[ + :, self.history_length - self.context_length :, ... + ] + sequence = past_target + control_sequence = past_control + sequence_length = self.history_length + subsequences_length = self.context_length + else: + time_feat = torch.cat( + ( + past_time_feat[:, self.history_length - self.context_length :, ...], + future_time_feat, + ), + dim=1, + ) + sequence = torch.cat((past_target, future_target), dim=1) + control_sequence = torch.cat((past_control, future_control), dim=1) + sequence_length = self.history_length + self.prediction_length + subsequences_length = self.context_length + self.prediction_length + + lags = self.get_lagged_subsequences( + sequence=sequence, + sequence_length=sequence_length, + indices=self.lags_seq, + subsequences_length=subsequences_length, + ) + + control_lags = self.get_lagged_subsequences( + sequence=control_sequence, + sequence_length=sequence_length, + indices=self.lags_seq, + subsequences_length=subsequences_length, + ) + + # scale is computed on the context length last units of the past target + # scale shape is (batch_size, 1, *target_shape) + _, scale = self.scaler( + past_target[:, self.context_length :, ...], + past_observed_values[:, self.context_length :, ...], + ) + + _, control_scale = self.control_scaler( + past_control[:, self.context_length :, ...], + past_observed_values[:, self.context_length :, ...], + ) + + # (batch_size, num_features) + embedded_cat = self.embedder(feat_static_cat) + + # in addition to embedding features, use the log scale as it can help + # prediction too + # (batch_size, num_features + prod(target_shape)) + static_feat = torch.cat( + ( + embedded_cat, + feat_static_real, + scale.log() if len(self.target_shape) == 0 else scale.squeeze(1).log(), + control_scale.log() + if len(self.target_shape) == 0 + else control_scale.squeeze(1).log(), + ), + dim=1, + ) + + # (batch_size, subsequences_length, num_features + 1) + repeated_static_feat = static_feat.unsqueeze(1).expand( + -1, subsequences_length, -1 + ) + + # (batch_size, sub_seq_len, *target_shape, num_lags) + lags_scaled = lags / scale.unsqueeze(-1) + control_lags_scaled = control_lags / control_scale.unsqueeze(-1) + + # from (batch_size, sub_seq_len, *target_shape, num_lags) + # to (batch_size, sub_seq_len, prod(target_shape) * num_lags) + input_lags = lags_scaled.reshape( + (-1, subsequences_length, len(self.lags_seq) * prod(self.target_shape)) + ) + + input_control_lags = control_lags_scaled.reshape( + (-1, subsequences_length, len(self.lags_seq) * prod(self.target_shape)) + ) + + # (batch_size, sub_seq_len, input_dim) + inputs = torch.cat( + (input_lags, input_control_lags, time_feat, repeated_static_feat), dim=-1 + ) + + # unroll encoder + outputs, state = self.rnn(inputs) + + # outputs: (batch_size, seq_len, num_cells) + # state: list of (num_layers, batch_size, num_cells) tensors + # scale: (batch_size, 1, *target_shape) + # static_feat: (batch_size, num_features + prod(target_shape)) + return outputs, state, scale, control_scale, static_feat + + +class CausalDeepARTrainingNetwork(CausalDeepARNetwork): + def distribution( + self, + feat_static_cat: torch.Tensor, + feat_static_real: torch.Tensor, + past_time_feat: torch.Tensor, + past_target: torch.Tensor, + past_observed_values: torch.Tensor, + future_time_feat: torch.Tensor, + future_target: torch.Tensor, + future_observed_values: torch.Tensor, + past_control: torch.Tensor, + future_control: torch.Tensor, + ) -> Tuple[Distribution]: + rnn_outputs, _, scale, control_scale, _ = self.unroll_encoder( + feat_static_cat=feat_static_cat, + feat_static_real=feat_static_real, + past_time_feat=past_time_feat, + past_target=past_target, + past_observed_values=past_observed_values, + future_time_feat=future_time_feat, + future_target=future_target, + past_control=past_control, + future_control=future_control, + ) + + control_dist_args = self.proj_control_args(rnn_outputs) + + control = torch.cat( + ( + past_control[:, self.history_length - self.context_length :, ...], + future_control, + ), + dim=1, + ) + + distr_args = self.proj_distr_args( + torch.cat((rnn_outputs, control.unsqueeze(-1)), dim=-1) + ) + + return ( + self.control_output.distribution(control_dist_args, scale=control_scale), + self.distr_output.distribution(distr_args, scale=scale), + ) + + def forward( + self, + feat_static_cat: torch.Tensor, + feat_static_real: torch.Tensor, + past_time_feat: torch.Tensor, + past_target: torch.Tensor, + past_observed_values: torch.Tensor, + future_time_feat: torch.Tensor, + future_target: torch.Tensor, + future_observed_values: torch.Tensor, + past_control: torch.Tensor, + future_control: torch.Tensor, + ) -> torch.Tensor: + + control_distr, distr = self.distribution( + feat_static_cat=feat_static_cat, + feat_static_real=feat_static_real, + past_time_feat=past_time_feat, + past_target=past_target, + past_observed_values=past_observed_values, + future_time_feat=future_time_feat, + future_target=future_target, + future_observed_values=future_observed_values, + past_control=past_control, + future_control=future_control, + ) + + # put together target sequence + # (batch_size, seq_len, *target_shape) + target = torch.cat( + ( + past_target[:, self.history_length - self.context_length :, ...], + future_target, + ), + dim=1, + ) + + control = torch.cat( + ( + past_control[:, self.history_length - self.context_length :, ...], + future_control, + ), + dim=1, + ) + + # (batch_size, seq_len) + loss = -control_distr.log_prob(control) - distr.log_prob(target) + + # (batch_size, seq_len, *target_shape) + observed_values = torch.cat( + ( + past_observed_values[ + :, self.history_length - self.context_length :, ... + ], + future_observed_values, + ), + dim=1, + ) + + # mask the loss at one time step iff one or more observations is missing in the target dimensions + # (batch_size, seq_len) + loss_weights = ( + observed_values + if (len(self.target_shape) == 0) + else observed_values.min(dim=-1, keepdim=False) + ) + + weighted_loss = weighted_average(loss, weights=loss_weights) + + return weighted_loss, loss + + +class CausalDeepARPredictionNetwork(CausalDeepARNetwork): + def __init__(self, num_parallel_samples: int = 100, **kwargs) -> None: + super().__init__(**kwargs) + self.num_parallel_samples = num_parallel_samples + + # for decoding the lags are shifted by one, at the first time-step + # of the decoder a lag of one corresponds to the last target value + self.shifted_lags = [l - 1 for l in self.lags_seq] + + def sampling_decoder( + self, + static_feat: torch.Tensor, + past_target: torch.Tensor, + past_control: torch.Tensor, + time_feat: torch.Tensor, + scale: torch.Tensor, + control_scale: torch.Tensor, + future_control: torch.Tensor, + begin_states: Union[torch.Tensor, List[torch.Tensor]], + ) -> torch.Tensor: + """ + Computes sample paths by unrolling the RNN starting with a initial + input and state. + + Parameters + ---------- + static_feat : Tensor + static features. Shape: (batch_size, num_static_features). + past_target : Tensor + target history. Shape: (batch_size, history_length). + time_feat : Tensor + time features. Shape: (batch_size, prediction_length, num_time_features). + scale : Tensor + tensor containing the scale of each element in the batch. Shape: (batch_size, 1, 1). + begin_states : List or Tensor + list of initial states for the LSTM layers or tensor for GRU. + the shape of each tensor of the list should be (num_layers, batch_size, num_cells) + Returns + -------- + Tensor + A tensor containing sampled paths. + Shape: (batch_size, num_sample_paths, prediction_length). + """ + + # blows-up the dimension of each tensor to batch_size * self.num_parallel_samples for increasing parallelism + repeated_past_target = past_target.repeat_interleave( + repeats=self.num_parallel_samples, dim=0 + ) + repeated_time_feat = time_feat.repeat_interleave( + repeats=self.num_parallel_samples, dim=0 + ) + repeated_past_control = past_control.repeat_interleave( + repeats=self.num_parallel_samples, dim=0 + ) + repeated_static_feat = static_feat.repeat_interleave( + repeats=self.num_parallel_samples, dim=0 + ).unsqueeze(1) + repeated_future_control = future_control.repeat_interleave( + repeats=self.num_parallel_samples, dim=0 + ) + repeated_scale = scale.repeat_interleave( + repeats=self.num_parallel_samples, dim=0 + ) + repeated_control_scale = control_scale.repeat_interleave( + repeats=self.num_parallel_samples, dim=0 + ) + + if self.cell_type == "LSTM": + repeated_states = [ + s.repeat_interleave(repeats=self.num_parallel_samples, dim=1) + for s in begin_states + ] + else: + repeated_states = begin_states.repeat_interleave( + repeats=self.num_parallel_samples, dim=1 + ) + + future_samples = [] + + # for each future time-units we draw new samples for this time-unit and update the state + for k in range(self.prediction_length): + # (batch_size * num_samples, 1, *target_shape, num_lags) + lags = self.get_lagged_subsequences( + sequence=repeated_past_target, + sequence_length=self.history_length + k, + indices=self.shifted_lags, + subsequences_length=1, + ) + control_lags = self.get_lagged_subsequences( + sequence=repeated_past_control, + sequence_length=self.history_length + k, + indices=self.shifted_lags, + subsequences_length=1, + ) + + # (batch_size * num_samples, 1, *target_shape, num_lags) + lags_scaled = lags / repeated_scale.unsqueeze(-1) + + control_lags_scaled = control_lags / repeated_control_scale.unsqueeze(-1) + + # from (batch_size * num_samples, 1, *target_shape, num_lags) + # to (batch_size * num_samples, 1, prod(target_shape) * num_lags) + input_lags = lags_scaled.reshape( + (-1, 1, prod(self.target_shape) * len(self.lags_seq)) + ) + + input_control_lags = control_lags_scaled.reshape( + (-1, 1, len(self.lags_seq) * prod(self.target_shape)) + ) + + # (batch_size * num_samples, 1, prod(target_shape) * num_lags + num_time_features + num_static_features) + decoder_input = torch.cat( + ( + input_lags, + input_control_lags, + repeated_time_feat[:, k : k + 1, :], + repeated_static_feat, + ), + dim=-1, + ) + + # output shape: (batch_size * num_samples, 1, num_cells) + # state shape: (batch_size * num_samples, num_cells) + rnn_outputs, repeated_states = self.rnn(decoder_input, repeated_states) + + control_dist_args = self.proj_control_args(rnn_outputs) + control_distr = self.control_output.distribution( + control_dist_args, repeated_control_scale + ) + + new_control_sample = control_distr.sample() + control = repeated_future_control[:, k : k + 1] + control[control != control] = new_control_sample[control != control] + repeated_past_control = torch.cat((repeated_past_control, control), dim=1) + + distr_args = self.proj_distr_args( + torch.cat((rnn_outputs, control.unsqueeze(-1)), dim=-1) + ) + + # compute likelihood of target given the predicted parameters + distr = self.distr_output.distribution(distr_args, scale=repeated_scale) + + # (batch_size * num_samples, 1, *target_shape) + new_samples = distr.sample() + + # (batch_size * num_samples, seq_len, *target_shape) + repeated_past_target = torch.cat((repeated_past_target, new_samples), dim=1) + future_samples.append(new_samples) + + # (batch_size * num_samples, prediction_length, *target_shape) + samples = torch.cat(future_samples, dim=1) + + # (batch_size, num_samples, prediction_length, *target_shape) + return samples.reshape( + ( + (-1, self.num_parallel_samples) + + (self.prediction_length,) + + self.target_shape + ) + ) + + # noinspection PyMethodOverriding,PyPep8Naming + def forward( + self, + feat_static_cat: torch.Tensor, # (batch_size, num_features) + feat_static_real: torch.Tensor, # (batch_size, num_features) + past_time_feat: torch.Tensor, # (batch_size, history_length, num_features) + past_target: torch.Tensor, # (batch_size, history_length, *target_shape) + past_observed_values: torch.Tensor, # (batch_size, history_length, *target_shape) + future_time_feat: torch.Tensor, # (batch_size, prediction_length, num_features) + past_control: torch.Tensor, + future_control: torch.Tensor, + ) -> torch.Tensor: + """ + Predicts samples, all tensors should have NTC layout. + Parameters + ---------- + feat_static_cat : (batch_size, num_features) + feat_static_real : (batch_size, num_features) + past_time_feat : (batch_size, history_length, num_features) + past_target : (batch_size, history_length, *target_shape) + past_observed_values : (batch_size, history_length, *target_shape) + future_time_feat : (batch_size, prediction_length, num_features) + + Returns + ------- + Tensor + Predicted samples + """ + + # unroll the decoder in "prediction mode", i.e. with past data only + _, state, scale, control_scale, static_feat = self.unroll_encoder( + feat_static_cat=feat_static_cat, + feat_static_real=feat_static_real, + past_time_feat=past_time_feat, + past_target=past_target, + past_observed_values=past_observed_values, + past_control=past_control, + future_time_feat=None, + future_target=None, + future_control=None, + ) + + return self.sampling_decoder( + past_target=past_target, + past_control=past_control, + time_feat=future_time_feat, + static_feat=static_feat, + future_control=future_control, + scale=scale, + control_scale=control_scale, + begin_states=state, + )