diff --git a/pts/model/deepvar/deepvar_estimator.py b/pts/model/deepvar/deepvar_estimator.py new file mode 100644 index 0000000..d63f8cf --- /dev/null +++ b/pts/model/deepvar/deepvar_estimator.py @@ -0,0 +1,248 @@ +from typing import List, Optional + + +import numpy as np +import pandas as pd +from pandas.tseries.frequencies import to_offset +import torch +import torch.nn as nn + +from pts import Trainer +from pts.model import PTSEstimator, PTSPredictor, copy_parameters +from pts.modules import DistributionOutput, LowRankMultivariateNormalOutput +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, +) + + +def get_lags_for_frequency(freq_str: str, num_lags: Optional[int] = None) -> List[int]: + offset = to_offset(freq_str) + multiple, granularity = offset.n, offset.name + + if granularity == "M": + lags = [[1, 12]] + elif granularity == "D": + lags = [[1, 7, 14]] + elif granularity == "B": + lags = [[1, 2]] + elif granularity == "H": + lags = [[1, 24, 168]] + elif granularity == "min": + lags = [[1, 4, 12, 24, 48]] + else: + lags = [[1]] + + # use less lags + output_lags = list([int(lag) for sub_list in lags for lag in sub_list]) + output_lags = sorted(list(set(output_lags))) + return output_lags[:num_lags] + + +class FourierDateFeatures(TimeFeature): + @validated() + def __init__(self, freq: str) -> None: + super().__init__() + # reoccurring freq + freqs = [ + "month", + "day", + "hour", + "minute", + "weekofyear", + "weekday", + "dayofweek", + "dayofyear", + "daysinmonth", + ] + + assert freq in freqs + self.freq = freq + + def __call__(self, index: pd.DatetimeIndex) -> np.ndarray: + values = getattr(index, self.freq) + num_values = max(values) + 1 + steps = [x * 2.0 * np.pi / num_values for x in values] + return np.vstack([np.cos(steps), np.sin(steps)]) + + +def time_features_from_frequency_str(freq_str: str) -> List[TimeFeature]: + offset = to_offset(freq_str) + multiple, granularity = offset.n, offset.name + + features = { + "M": ["weekofyear"], + "W": ["daysinmonth", "weekofyear"], + "D": ["dayofweek"], + "B": ["dayofweek", "dayofyear"], + "H": ["hour", "dayofweek"], + "min": ["minute", "hour", "dayofweek"], + "T": ["minute", "hour", "dayofweek"], + } + + assert granularity in features, f"freq {granularity} not supported" + + feature_classes: List[TimeFeature] = [ + FourierDateFeatures(freq=freq) for freq in features[granularity] + ] + return feature_classes + + +class DeepVAREstimator(PTSEstimator): + def __init__( + self, + 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, + distr_output: Optional[DistributionOutput] = None, + rank: Optional[int] = 5, + scaling: bool = True, + pick_incomplete: bool = False, + lags_seq: Optional[List[int]] = None, + time_features: Optional[List[TimeFeature]] = None, + conditioning_length: int = 200, + use_marginal_transformation=False, + **kwargs, + ) -> None: + super().__init__(trainer=trainer, **kwargs) + + self.freq = freq + self.context_length = ( + context_length if context_length is not None else prediction_length + ) + + if distr_output is not None: + self.distr_output = distr_output + else: + self.distr_output = LowRankMultivariateNormalOutput( + dim=target_dim, rank=rank + ) + + 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.conditioning_length = conditioning_length + self.use_marginal_t + + self.lags_seq = ( + lags_seq if lags_seq is not None else get_lags_for_frequency(freq_str=freq) + ) + + 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.pick_incomplete = pick_incomplete + self.scaling = scaling + + if self.use_marginal_transformation: + self.output_transform: Optional[ + Callable + ] = cdf_to_gaussian_forward_transform + else: + self.output_transform = None + + def create_transformation(self) -> Transformation: + def use_marginal_transformation( + marginal_transformation: bool + ) -> Transformation: + if marginal_transformation: + return CDFtoGaussianTransform( + target_field=FieldName.TARGET, + observed_values_field=FieldName.OBSERVED_VALUES, + max_context_length=self.conditioning_length, + target_dim=self.target_dim, + ) + else: + return RenameFields( + { + f"past_{FieldName.TARGET}": f"past_{FieldName.TARGET}_cdf", + f"future_{FieldName.TARGET}": f"future_{FieldName.TARGET}_cdf", + } + ) + + 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.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, + ), + use_marginal_transformation(self.use_marginal_transformation), + ] + ) + + \ No newline at end of file