From 52fbc77f2127a422c66b83a1ca978843edd0b664 Mon Sep 17 00:00:00 2001 From: "Dr. Kashif Rasul" Date: Mon, 13 Jan 2020 11:28:31 +0100 Subject: [PATCH] added tempFlow model --- pts/model/tempflow/__init__.py | 2 + pts/model/tempflow/tempflow_estimator.py | 222 +++++++++ pts/model/tempflow/tempflow_network.py | 580 +++++++++++++++++++++++ pts/modules/__init__.py | 1 + pts/modules/distribution_output.py | 12 +- 5 files changed, 808 insertions(+), 9 deletions(-) create mode 100644 pts/model/tempflow/__init__.py create mode 100644 pts/model/tempflow/tempflow_estimator.py create mode 100644 pts/model/tempflow/tempflow_network.py diff --git a/pts/model/tempflow/__init__.py b/pts/model/tempflow/__init__.py new file mode 100644 index 0000000..270ae57 --- /dev/null +++ b/pts/model/tempflow/__init__.py @@ -0,0 +1,2 @@ +from .tempflow_estimator import TempFlowEstimator +from .tempflow_network import TempFlowTrainingNetwork, TempFlowPredictionNetwork diff --git a/pts/model/tempflow/tempflow_estimator.py b/pts/model/tempflow/tempflow_estimator.py new file mode 100644 index 0000000..a5b2c57 --- /dev/null +++ b/pts/model/tempflow/tempflow_estimator.py @@ -0,0 +1,222 @@ +from typing import List, Optional + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn + +from pts import Trainer +from pts.model import PTSEstimator, PTSPredictor, copy_parameters +from pts.modules import RealNVP +from pts.dataset import FieldName +from pts.transform import ( + Transformation, + Chain, + InstanceSplitter, + ExpectedNumInstanceSampler, + CDFtoGaussianTransform, + cdf_to_gaussian_forward_transform, + RenameFields, + AsNumpyArray, + ExpandDimArray, + AddObservedValuesIndicator, + AddTimeFeatures, + VstackFeatures, + SetFieldIfNotPresent, + TargetDimIndicator, +) +from pts.feature import ( + TimeFeature, + fourier_time_features_from_frequency_str, + get_fourier_lags_for_frequency, +) +from + +from .tempflow_network import TempFlowTrainingNetwork, TempFlowPredictionNetwork + + +class TempFlowEstimator(PTSEstimator): + def __init__( + self, + input_size: int, + freq: str, + prediction_length: int, + target_dim: int, + trainer: Trainer = Trainer(), + context_length: Optional[int] = None, + num_layers: int = 2, + num_cells: int = 40, + cell_type: str = "LSTM", + num_parallel_samples: int = 100, + dropout_rate: float = 0.1, + + cardinality: List[int] = [1], + embedding_dimension: int = 5, + + flow_type = "RealNVP", + n_blocks = 3, + hidden_size = 100, + n_hidden = 2, + conditioning_length: int = 200, + + scaling: bool = True, + pick_incomplete: bool = False, + lags_seq: Optional[List[int]] = None, + time_features: Optional[List[TimeFeature]] = None, + + **kwargs, + ) -> None: + super().__init__(trainer=trainer, **kwargs) + + self.freq = freq + self.context_length = ( + context_length if context_length is not None else prediction_length + ) + + self.input_size = input_size + self.prediction_length = prediction_length + self.target_dim = target_dim + self.num_layers = num_layers + self.num_cells = num_cells + self.cell_type = cell_type + self.num_parallel_samples = num_parallel_samples + self.dropout_rate = dropout_rate + self.cardinality = cardinality + self.embedding_dimension = embedding_dimension + + self.flow_type = flow_type + self.n_blocks = n_blocks + self.hidden_size = hidden_size + self.n_hidden = n_hidden + self.conditioning_length = conditioning_length + + self.lags_seq = ( + lags_seq + if lags_seq is not None + else get_fourier_lags_for_frequency(freq_str=freq) + ) + + self.time_features = ( + time_features + if time_features is not None + else fourier_time_features_from_frequency_str(self.freq) + ) + + self.history_length = self.context_length + max(self.lags_seq) + self.pick_incomplete = pick_incomplete + self.scaling = scaling + + def create_transformation(self) -> Transformation: + return Chain( + [ + AsNumpyArray( + field=FieldName.TARGET, + expected_ndim=1 + len(self.distr_output.event_shape), + ), + # maps the target to (1, T) + # if the target data is uni dimensional + ExpandDimArray( + field=FieldName.TARGET, + axis=0 if self.distr_output.event_shape[0] == 1 else None, + ), + AddObservedValuesIndicator( + target_field=FieldName.TARGET, + output_field=FieldName.OBSERVED_VALUES, + ), + AddTimeFeatures( + start_field=FieldName.START, + target_field=FieldName.TARGET, + output_field=FieldName.FEAT_TIME, + time_features=self.time_features, + pred_length=self.prediction_length, + ), + VstackFeatures( + output_field=FieldName.FEAT_TIME, + input_fields=[FieldName.FEAT_TIME], + ), + SetFieldIfNotPresent(field=FieldName.FEAT_STATIC_CAT, value=[0]), + TargetDimIndicator( + field_name="target_dimension_indicator", + target_field=FieldName.TARGET, + ), + AsNumpyArray(field=FieldName.FEAT_STATIC_CAT, expected_ndim=1), + InstanceSplitter( + target_field=FieldName.TARGET, + is_pad_field=FieldName.IS_PAD, + start_field=FieldName.START, + forecast_start_field=FieldName.FORECAST_START, + train_sampler=ExpectedNumInstanceSampler(num_instances=1), + past_length=self.history_length, + future_length=self.prediction_length, + time_series_fields=[ + FieldName.FEAT_TIME, + FieldName.OBSERVED_VALUES, + ], + pick_incomplete=self.pick_incomplete, + ), + ] + ) + + def create_training_network(self, device: torch.device) -> TempFlowTrainingNetwork: + return TempFlowTrainingNetwork( + input_size=self.input_size, + target_dim=self.target_dim, + num_parallel_samples=self.num_parallel_samples, + 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, + dropout_rate=self.dropout_rate, + cardinality=self.cardinality, + embedding_dimension=self.embedding_dimension, + lags_seq=self.lags_seq, + scaling=self.scaling, + flow_type=self.flow_type, + n_blocks=self.n_blocks, + hidden_size = self.hidden_size, + n_hidden = self.n_hidden, + conditioning_length=self.conditioning_length, + ).to(device) + + def create_predictor( + self, + transformation: Transformation, + trained_network: TempFlowTrainingNetwork, + device: torch.device, + ) -> PTSPredictor: + prediction_network = TempFlowPredictionNetwork( + input_size=self.input_size, + target_dim=self.target_dim, + num_parallel_samples=self.num_parallel_samples, + 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, + dropout_rate=self.dropout_rate, + cardinality=self.cardinality, + embedding_dimension=self.embedding_dimension, + lags_seq=self.lags_seq, + scaling=self.scaling, + flow_type=self.flow_type, + n_blocks=self.n_blocks, + hidden_size = self.hidden_size, + n_hidden = self.n_hidden, + conditioning_length=self.conditioning_length, + num_parallel_samples=self.num_parallel_samples + ).to(device) + + copy_parameters(trained_network, prediction_network) + + return PTSPredictor( + input_transform=transformation, + prediction_net=prediction_network, + batch_size=self.trainer.batch_size, + freq=self.freq, + prediction_length=self.prediction_length, + device=device, + output_transform=self.output_transform, + ) diff --git a/pts/model/tempflow/tempflow_network.py b/pts/model/tempflow/tempflow_network.py new file mode 100644 index 0000000..72e8b96 --- /dev/null +++ b/pts/model/tempflow/tempflow_network.py @@ -0,0 +1,580 @@ +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn + +import numpy as np + +from pts.modules import RealNVP, FlowOutput, MeanScaler, NOPScaler +from pts.model import weighted_average + + +class TempFlowTrainingNetwork(nn.Module): + 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, + dropout_rate: float, + lags_seq: List[int], + target_dim: int, + conditioning_length: int, + flow_type: str, + n_blocks: int, + hidden_size: int, + n_hidden: int, + cardinality: List[int] = [1], + embedding_dimension: int = 1, + scaling: bool = True, + **kwargs, + ) -> None: + super().__init__(**kwargs) + + assert len(set(lags_seq)) == len(lags_seq), "no duplicated lags allowed!" + lags_seq.sort() + self.lags_seq = lags_seq + + rnn_cls = {"LSTM": nn.LSTM, "GRU": nn.GRU}[self.cell_type] + self.rnn = rnn_cls( + input_size=input_size, + hidden_size=num_cells, + num_layers=num_layers, + dropout=dropout_rate, + batch_first=True, + ) + + flow_cls = {"RealNVP": RealNVP}[self.flow_type] + self.flow = flow_cls( + input_size=input_size, + n_blocks=n_blocks, + n_hidden=n_hidden, + hidden_size=hidden_size, + cond_label_size=conditioning_length, + ) + + self.distr_output = FlowOutput( + self.flow, input_size=input_size, cond_size=conditioning_length + ) + + self.proj_dist_args = distr_output.get_args_proj(num_cells) + + self.embed_dim = 1 + self.embed = nn.Embedding( + num_embeddings=self.target_dim, embedding_dim=self.embed_dim + ) + + if scaling: + self.scaler = MeanScaler(keepdim=True) + else: + self.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 + the sequence from which lagged subsequences should be extracted. + Shape: (N, T, C). + sequence_length + length of sequence in the T (time) dimension (axis = 1). + indices + list of lag indices to be used. + subsequences_length + 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, :]. + """ + # we must have: history_length + begin_index >= 0 + # that is: history_length - lag_index - sequence_length >= 0 + # hence the following assert + assert max(indices) + subsequences_length <= sequence_length, ( + f"lags cannot go further than history length, found lag " + f"{max(indices)} 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, ...].unsqueeze(1)) + return torch.cat(lagged_values, dim=1).permute(0, 2, 3, 1) + + def unroll( + self, + lags: torch.Tensor, + scale: torch.Tensor, + time_feat: torch.Tensor, + target_dimension_indicator: torch.Tensor, + unroll_length: int, + begin_state: Optional[Union[List[torch.Tensor], torch.Tensor]] = None, + ) -> Tuple[ + torch.Tensor, + Union[List[torch.Tensor], torch.Tensor], + torch.Tensor, + torch.Tensor, + ]: + + # (batch_size, sub_seq_len, target_dim, num_lags) + lags_scaled = lags / scale.unsqueeze(-1) + + # assert_shape( + # lags_scaled, (-1, unroll_length, self.target_dim, len(self.lags_seq)), + # ) + + input_lags = lags_scaled.reshape( + (-1, unroll_length, len(self.lags_seq) * self.target_dim) + ) + + # (batch_size, target_dim, embed_dim) + index_embeddings = self.embed(target_dimension_indicator) + # assert_shape(index_embeddings, (-1, self.target_dim, self.embed_dim)) + + # (batch_size, seq_len, target_dim * embed_dim) + repeated_index_embeddings = ( + index_embeddings.unsqueeze(1) + .expand(-1, unroll_length, -1, -1) + .reshape((-1, unroll_length, self.target_dim * self.embed_dim)) + ) + + # (batch_size, sub_seq_len, input_dim) + inputs = torch.cat((input_lags, repeated_index_embeddings, time_feat), dim=-1) + + # unroll encoder + outputs, state = self.rnn(inputs, begin_state) + + # assert_shape(outputs, (-1, unroll_length, self.num_cells)) + # for s in state: + # assert_shape(s, (-1, self.num_cells)) + + # assert_shape( + # lags_scaled, (-1, unroll_length, self.target_dim, len(self.lags_seq)), + # ) + + return outputs, state, lags_scaled, inputs + + def unroll_encoder( + self, + past_time_feat: torch.Tensor, + past_target_cdf: torch.Tensor, + past_observed_values: torch.Tensor, + past_is_pad: torch.Tensor, + future_time_feat: Optional[torch.Tensor], + future_target_cdf: Optional[torch.Tensor], + target_dimension_indicator: torch.Tensor, + ) -> Tuple[ + torch.Tensor, + Union[List[torch.Tensor], torch.Tensor], + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + """ + Unrolls the RNN encoder over past and, if present, future data. + Returns outputs and state of the encoder, plus the scale of + past_target_cdf and a vector of static features that was constructed + and fed as input to the encoder. All tensor arguments should have NTC + layout. + + Parameters + ---------- + past_time_feat + Past time features (batch_size, history_length, num_features) + past_target_cdf + Past marginal CDF transformed target values (batch_size, + history_length, target_dim) + past_observed_values + Indicator whether or not the values were observed (batch_size, + history_length, target_dim) + past_is_pad + Indicator whether the past target values have been padded + (batch_size, history_length) + future_time_feat + Future time features (batch_size, prediction_length, num_features) + future_target_cdf + Future marginal CDF transformed target values (batch_size, + prediction_length, target_dim) + target_dimension_indicator + Dimensionality of the time series (batch_size, target_dim) + + Returns + ------- + outputs + RNN outputs (batch_size, seq_len, num_cells) + states + RNN states. Nested list with (batch_size, num_cells) tensors with + dimensions target_dim x num_layers x (batch_size, num_cells) + scale + Mean scales for the time series (batch_size, 1, target_dim) + lags_scaled + Scaled lags(batch_size, sub_seq_len, target_dim, num_lags) + inputs + inputs to the RNN + + """ + + past_observed_values = torch.min( + past_observed_values, 1 - past_is_pad.unsqueeze(-1) + ) + + if future_time_feat is None or future_target_cdf is None: + time_feat = past_time_feat[:, -self.context_length :, ...] + sequence = past_target_cdf + sequence_length = self.history_length + subsequences_length = self.context_length + else: + time_feat = torch.cat( + (past_time_feat[:, -self.context_length :, ...], future_time_feat), + dim=1, + ) + sequence = torch.cat((past_target_cdf, future_target_cdf), dim=1) + sequence_length = self.history_length + self.prediction_length + subsequences_length = self.context_length + self.prediction_length + + # (batch_size, sub_seq_len, target_dim, num_lags) + lags = self.get_lagged_subsequences( + sequence=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_dim) + _, scale = self.scaler( + past_target_cdf[:, -self.context_length :, ...], + past_observed_values[:, -self.context_length :, ...], + ) + + # import pdb; pdb.set_trace() + + outputs, states, lags_scaled, inputs = self.unroll( + lags=lags, + scale=scale, + time_feat=time_feat, + target_dimension_indicator=target_dimension_indicator, + unroll_length=subsequences_length, + begin_state=None, + ) + + return outputs, states, scale, lags_scaled, inputs + + def distr( + self, rnn_outputs: torch.Tensor, scale: torch.Tensor, + ): + """ + Returns the distribution of DeepVAR with respect to the RNN outputs. + + Parameters + ---------- + rnn_outputs + Outputs of the unrolled RNN (batch_size, seq_len, num_cells) + scale + Mean scale for each time series (batch_size, 1, target_dim) + + Returns + ------- + distr + Distribution instance + distr_args + Distribution arguments + """ + distr_args = self.proj_dist_args(rnn_outputs) + + # compute likelihood of target given the predicted parameters + distr = self.distr_output.distribution(distr_args, scale=scale) + + return distr, distr_args + + def forward( + self, + target_dimension_indicator: torch.Tensor, + past_time_feat: torch.Tensor, + past_target_cdf: torch.Tensor, + past_observed_values: torch.Tensor, + past_is_pad: torch.Tensor, + future_time_feat: torch.Tensor, + future_target_cdf: torch.Tensor, + future_observed_values: torch.Tensor, + ) -> Tuple[torch.Tensor, ...]: + """ + Computes the loss for training DeepVAR, all inputs tensors representing + time series have NTC layout. + + Parameters + ---------- + target_dimension_indicator + Indices of the target dimension (batch_size, target_dim) + past_time_feat + Dynamic features of past time series (batch_size, history_length, + num_features) + past_target_cdf + Past marginal CDF transformed target values (batch_size, + history_length, target_dim) + past_observed_values + Indicator whether or not the values were observed (batch_size, + history_length, target_dim) + past_is_pad + Indicator whether the past target values have been padded + (batch_size, history_length) + future_time_feat + Future time features (batch_size, prediction_length, num_features) + future_target_cdf + Future marginal CDF transformed target values (batch_size, + prediction_length, target_dim) + future_observed_values + Indicator whether or not the future values were observed + (batch_size, prediction_length, target_dim) + + Returns + ------- + distr + Loss with shape (batch_size, 1) + likelihoods + Likelihoods for each time step + (batch_size, context + prediction_length, 1) + distr_args + Distribution arguments (context + prediction_length, + number_of_arguments) + """ + + seq_len = self.context_length + self.prediction_length + + # unroll the decoder in "training mode", i.e. by providing future data + # as well + rnn_outputs, _, scale, _, inputs = self.unroll_encoder( + past_time_feat=past_time_feat, + past_target_cdf=past_target_cdf, + past_observed_values=past_observed_values, + past_is_pad=past_is_pad, + future_time_feat=future_time_feat, + future_target_cdf=future_target_cdf, + target_dimension_indicator=target_dimension_indicator, + ) + + # put together target sequence + # (batch_size, seq_len, target_dim) + target = torch.cat( + (past_target_cdf[:, -self.context_length :, ...], future_target_cdf), dim=1, + ) + + # assert_shape(target, (-1, seq_len, self.target_dim)) + + distr, distr_args = self.distr(rnn_outputs=rnn_outputs, scale=scale) + + # we sum the last axis to have the same shape for all likelihoods + # (batch_size, subseq_length, 1) + likelihoods = -distr.log_prob(target).unsqueeze(-1) + + # import pdb; pdb.set_trace() + + # assert_shape(likelihoods, (-1, seq_len, 1)) + + past_observed_values = torch.min( + past_observed_values, 1 - past_is_pad.unsqueeze(-1) + ) + + # (batch_size, subseq_length, target_dim) + observed_values = torch.cat( + ( + past_observed_values[:, -self.context_length :, ...], + future_observed_values, + ), + dim=1, + ) + + # mask the loss at one time step if one or more observations is missing + # in the target dimensions (batch_size, subseq_length, 1) + loss_weights, _ = observed_values.min(dim=-1, keepdim=True) + + # assert_shape(loss_weights, (-1, seq_len, 1)) + + loss = weighted_average(likelihoods, weights=loss_weights, dim=1) + + # assert_shape(loss, (-1, -1, 1)) + + # self.distribution = distr + + return (loss.mean(), likelihoods) + distr_args + + +class TempFlowPredictionNetwork(TempFlowTrainingNetwork): + def __init__(self, num_parallel_samples: int, **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, + past_target_cdf: torch.Tensor, + target_dimension_indicator: torch.Tensor, + time_feat: torch.Tensor, + scale: torch.Tensor, + begin_states: Union[List[torch.Tensor], torch.Tensor] + ) -> torch.Tensor: + """ + Computes sample paths by unrolling the RNN starting with a initial + input and state. + + Parameters + ---------- + past_target_cdf + Past marginal CDF transformed target values (batch_size, + history_length, target_dim) + target_dimension_indicator + Indices of the target dimension (batch_size, target_dim) + time_feat + Dynamic features of future time series (batch_size, history_length, + num_features) + scale + Mean scale for each time series (batch_size, 1, target_dim) + begin_states + List of initial states for the RNN layers (batch_size, num_cells) + Returns + -------- + sample_paths : Tensor + A tensor containing sampled paths. Shape: (1, num_sample_paths, + prediction_length, target_dim). + """ + + def repeat(tensor, dim=0): + return tensor.repeat_interleave(repeats=self.num_parallel_samples, dim=dim) + + # blows-up the dimension of each tensor to + # batch_size * self.num_sample_paths for increasing parallelism + repeated_past_target_cdf = repeat(past_target_cdf) + repeated_time_feat = repeat(time_feat) + repeated_scale = repeat(scale) + repeated_target_dimension_indicator = repeat(target_dimension_indicator) + + if self.cell_type == "LSTM": + repeated_states = [repeat(s, dim=1) for s in begin_states] + else: + repeated_states = repeat(begin_states, 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): + lags = self.get_lagged_subsequences( + sequence=repeated_past_target_cdf, + sequence_length=self.history_length + k, + indices=self.shifted_lags, + subsequences_length=1, + ) + + rnn_outputs, repeated_states, _, _ = self.unroll( + begin_state=repeated_states, + lags=lags, + scale=repeated_scale, + time_feat=repeated_time_feat[:, k : k + 1, ...], + target_dimension_indicator=repeated_target_dimension_indicator, + unroll_length=1, + ) + + distr, distr_args = self.distr( + rnn_outputs=rnn_outputs, scale=repeated_scale, + ) + + # (batch_size, 1, target_dim) + new_samples = distr.sample() + + # (batch_size, seq_len, target_dim) + future_samples.append(new_samples) + repeated_past_target_cdf = torch.cat( + (repeated_past_target_cdf, new_samples), dim=1 + ) + + # (batch_size * num_samples, prediction_length, target_dim) + samples = torch.cat(future_samples, dim=1) + + # (batch_size, num_samples, prediction_length, target_dim) + return samples.reshape( + (-1, self.num_parallel_samples, self.prediction_length, self.target_dim,) + ) + + def forward( + self, + target_dimension_indicator: torch.Tensor, + past_time_feat: torch.Tensor, + past_target_cdf: torch.Tensor, + past_observed_values: torch.Tensor, + past_is_pad: torch.Tensor, + future_time_feat: torch.Tensor, + ) -> torch.Tensor: + """ + Predicts samples given the trained DeepVAR model. + All tensors should have NTC layout. + Parameters + ---------- + target_dimension_indicator + Indices of the target dimension (batch_size, target_dim) + past_time_feat + Dynamic features of past time series (batch_size, history_length, + num_features) + past_target_cdf + Past marginal CDF transformed target values (batch_size, + history_length, target_dim) + past_observed_values + Indicator whether or not the values were observed (batch_size, + history_length, target_dim) + past_is_pad + Indicator whether the past target values have been padded + (batch_size, history_length) + future_time_feat + Future time features (batch_size, prediction_length, num_features) + + Returns + ------- + sample_paths : Tensor + A tensor containing sampled paths (1, num_sample_paths, + prediction_length, target_dim). + + """ + + # mark padded data as unobserved + # (batch_size, target_dim, seq_len) + past_observed_values = torch.min( + past_observed_values, 1 - past_is_pad.unsqueeze(-1) + ) + + # unroll the decoder in "prediction mode", i.e. with past data only + _, state, scale, _, _ = self.unroll_encoder( + past_time_feat=past_time_feat, + past_target_cdf=past_target_cdf, + past_observed_values=past_observed_values, + past_is_pad=past_is_pad, + future_time_feat=None, + future_target_cdf=None, + target_dimension_indicator=target_dimension_indicator, + ) + + return self.sampling_decoder( + past_target_cdf=past_target_cdf, + target_dimension_indicator=target_dimension_indicator, + time_feat=future_time_feat, + scale=scale, + begin_states=state, + ) diff --git a/pts/modules/__init__.py b/pts/modules/__init__.py index 49afc7b..6911fb9 100644 --- a/pts/modules/__init__.py +++ b/pts/modules/__init__.py @@ -8,6 +8,7 @@ from .distribution_output import ( IndependentNormalOutput, LowRankMultivariateNormalOutput, MultivariateNormalOutput, + FlowOutput, ) from .lambda_layer import LambdaLayer from .feature import FeatureEmbedder, FeatureAssembler diff --git a/pts/modules/distribution_output.py b/pts/modules/distribution_output.py index b5953b5..a5e6369 100644 --- a/pts/modules/distribution_output.py +++ b/pts/modules/distribution_output.py @@ -240,17 +240,11 @@ class MultivariateNormalOutput(DistributionOutput): -class RealNVPOutput(DistributionOutput): - def __init__(self, input_size, cond_size, n_blocks=3, n_hidden=1, hidden_size=100): +class FlowOutput(DistributionOutput): + def __init__(self, flow, input_size, cond_size): self.args_dim = {"cond": cond_size} + self.flow = flow self.dim = input_size - self.flow = RealNVP( - n_blocks=n_blocks, - input_size=input_size, - hidden_size=hidden_size, - n_hidden=n_hidden, - cond_label_size=cond_size - ) def domain_map(self, cond): return (cond,)