diff --git a/pts/model/deepar/deepar_estimator.py b/pts/model/deepar/deepar_estimator.py index a3b602c..bf89313 100644 --- a/pts/model/deepar/deepar_estimator.py +++ b/pts/model/deepar/deepar_estimator.py @@ -7,10 +7,16 @@ from pts.feature import ( TimeFeature, get_lags_for_frequency, time_features_from_frequency_str, + Transformation, Chain, + RemoveFields, SetField, AsNumpyArray, + AddObservedValuesIndicator, AddTimeFeatures, + AddAgeFeature, VstackFeatures, InstanceSplitter, ) +from pts.dataset import FieldName, ExpectedNumInstanceSampler from pts.model import PTSEstimator from pts.modules import DistributionOutput, StudentTOutput +from .deepar_network import DeepARTrainingNetwork class DeepAREstimator(PTSEstimator): def __init__( @@ -70,3 +76,105 @@ class DeepAREstimator(PTSEstimator): self.history_length = self.context_length + max(self.lags_seq) self.num_parallel_samples = num_parallel_samples + + def create_transformation(self) -> Transformation: + remove_field_names = [FieldName.FEAT_DYNAMIC_CAT] + 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) + + return Chain( + [RemoveFields(field_names=remove_field_names)] + + ( + [SetField(output_field=FieldName.FEAT_STATIC_CAT, value=[0.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=self.dtype, + ), + 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, + ), + 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 [] + ), + ), + 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, + ], + ), + ] + ) + + def create_training_network(self) -> DeepARTrainingNetwork: + return DeepARTrainingNetwork( + 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, + dropout_rate=self.dropout_rate, + cardinality=self.cardinality, + embedding_dimension=self.embedding_dimension, + lags_seq=self.lags_seq, + scaling=self.scaling, + dtype=self.dtype, + ) + \ No newline at end of file diff --git a/pts/model/deepar/deepar_network.py b/pts/model/deepar/deepar_network.py index f8f60ee..bac67d6 100644 --- a/pts/model/deepar/deepar_network.py +++ b/pts/model/deepar/deepar_network.py @@ -4,3 +4,6 @@ import torch.nn as nn class DeepARNetwork(nn.Module): pass + +class DeepARTrainingNetwork(DeepARNetwork): + pass \ No newline at end of file