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 .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=2,), # maps the target to (1, T) # if the target data is uni dimensional ExpandDimArray(field=FieldName.TARGET, axis=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, ), RenameFields( { f"past_{FieldName.TARGET}": f"past_{FieldName.TARGET}_cdf", f"future_{FieldName.TARGET}": f"future_{FieldName.TARGET}_cdf", } ), ] ) def create_training_network(self, device: torch.device) -> TempFlowTrainingNetwork: return TempFlowTrainingNetwork( input_size=self.input_size, target_dim=self.target_dim, 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_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=None, )