initial model

pts/dataset/loader.py

@@ -15,13 +15,13 @@ DataBatch = Dict[str, Any]
 
 class BatchBuffer:
     def __init__(
-        self, batch_size: int, device: torch.device, float_type: np.dtype = np.float32
+        self, batch_size: int, device: torch.device, dtype: np.dtype = np.float32
     ) -> None:
         self._buffers: Dict[Any, List[Any]] = defaultdict(list)
         self.batch_size = batch_size
         self._size = 0
         self.device = device
-        self.float_type = float_type
+        self.dtype = dtype
 
     def add(self, d: Dict[str, List[np.ndarray]]):
         if self._buffers:
@@ -46,7 +46,7 @@ class BatchBuffer:
         if isinstance(xs[0], np.ndarray):
             data = np.asarray(xs)
             if data.dtype.kind == "f":
-                data = data.astype(self.float_type)
+                data = data.astype(self.dtype)
             return torch.from_numpy(data).to(device=self.device, non_blocking=True)
         elif isinstance(xs[0], torch.Tensor):
             return torch.stack(*xs)
@@ -75,7 +75,7 @@ class DataLoader(Iterable[DataEntry]):
         The size of the batches to emit.
     device
         device to use to store data on.
-    float_type
+    dtype
         Floating point type to use.
     """
 
@@ -85,13 +85,13 @@ class DataLoader(Iterable[DataEntry]):
         transform: Transformation,
         batch_size: int,
         device: torch.device,
-        float_type: np.dtype = np.float32,
+        dtype: np.dtype = np.float32,
     ) -> None:
         self.dataset = dataset
         self.transform = transform
         self.batch_size = batch_size
         self.device = device
-        self.float_type = float_type
+        self.dtype = dtype
 
 
 class TrainDataLoader(DataLoader):
@@ -114,7 +114,7 @@ class TrainDataLoader(DataLoader):
         device to use to store data on.
     num_batches_per_epoch
         Number of batches to return in one complete iteration over this object.
-    float_type
+    dtype
         Floating point type to use.
     """
 
@@ -125,18 +125,18 @@ class TrainDataLoader(DataLoader):
         batch_size: int,
         device: torch.device,
         num_batches_per_epoch: int,
-        float_type: np.dtype = np.float32,
+        dtype: np.dtype = np.float32,
         shuffle_for_training: bool = True,
         num_batches_for_shuffling: int = 10,
     ) -> None:
-        super().__init__(dataset, transform, batch_size, device, float_type)
+        super().__init__(dataset, transform, batch_size, device, dtype)
         self.num_batches_per_epoch = num_batches_per_epoch
         self.shuffle_for_training = shuffle_for_training
         self._num_buffered_batches = (
             num_batches_for_shuffling if shuffle_for_training else 1
         )
         self._cur_iter: Optional[Iterator] = None
-        self._buffer = BatchBuffer(self.batch_size, device, float_type)
+        self._buffer = BatchBuffer(self.batch_size, device, dtype)
 
     def _emit_batches_while_buffer_larger_than(self, thresh) -> Iterator[DataBatch]:
         if self.shuffle_for_training:
@@ -196,12 +196,12 @@ class InferenceDataLoader(DataLoader):
         The size of the batches to emit.
     device
         device to use to store data on.
-    float_type
+    dtype
         Floating point type to use.
     """
 
     def __iter__(self) -> Iterator[DataBatch]:
-        buffer = BatchBuffer(self.batch_size, self.device, self.float_type)
+        buffer = BatchBuffer(self.batch_size, self.device, self.dtype)
         for data_entry in self.transform(iter(self.dataset), is_train=False):
             buffer.add(data_entry)
             if len(buffer) >= self.batch_size:

pts/feature/transform.py

@@ -408,6 +408,7 @@ class AddObservedValuesIndicator(SimpleTransformation):
         If set to true (default) missing values will be replaced. Otherwise
         they will not be replaced. In any case the indicator is included in the
         result.
+    dtype
     """
     def __init__(
             self,
@@ -415,11 +416,13 @@ class AddObservedValuesIndicator(SimpleTransformation):
             output_field: str,
             dummy_value: int = 0,
             convert_nans: bool = True,
+            dtype: np.dtype = np.float32,
     ) -> None:
         self.dummy_value = dummy_value
         self.target_field = target_field
         self.output_field = output_field
         self.convert_nans = convert_nans
+        self.dtype = dtype
 
     def transform(self, data: DataEntry) -> DataEntry:
         value = data[self.target_field]
@@ -430,8 +433,8 @@ class AddObservedValuesIndicator(SimpleTransformation):
             value[nan_indices] = self.dummy_value
 
         data[self.target_field] = value
-        # Invert bool array so that missing values are zeros and store as float
-        data[self.output_field] = np.invert(nan_entries).astype(np.float32)
+        # Invert bool array so that missing values are zeros and store as dtype
+        data[self.output_field] = np.invert(nan_entries).astype(self.dtype)
         return data
 
 
@@ -588,6 +591,7 @@ class AddAgeFeature(MapTransformation):
         Prediction length
     log_scale
         If set to true the age feature grows logarithmically otherwise linearly over time.
+    dtype
     """
     def __init__(
             self,
@@ -595,12 +599,14 @@ class AddAgeFeature(MapTransformation):
             output_field: str,
             pred_length: int,
             log_scale: bool = True,
+            dtype: np.dtype = np.float32,
     ) -> None:
         self.pred_length = pred_length
         self.target_field = target_field
         self.feature_name = output_field
         self.log_scale = log_scale
         self._age_feature = np.zeros(0)
+        self.dtype = dtype
 
     def map_transform(self, data: DataEntry, is_train: bool) -> DataEntry:
         length = target_transformation_length(data[self.target_field],
@@ -608,9 +614,9 @@ class AddAgeFeature(MapTransformation):
                                               is_train=is_train)
 
         if self.log_scale:
-            age = np.log10(2.0 + np.arange(length, dtype=np.float32))
+            age = np.log10(2.0 + np.arange(length, dtype=self.dtype))
         else:
-            age = np.arange(length, dtype=np.float32)
+            age = np.arange(length, dtype=self.dtype)
 
         data[self.feature_name] = age.reshape((1, length))
 

pts/feature/utils.py

@@ -18,4 +18,4 @@ def get_granularity(freq_str: str) -> Tuple[int, str]:
     groups = m.groups()
     multiple = int(groups[1]) if groups[1] is not None else 1
     granularity = groups[2]
-    return multiple, granularity
+    return multiple, granularity

pts/model/__init__.py

@@ -1,3 +1,4 @@
-from pts.model.estimator import Estimator
+from pts.model.estimator import Estimator, PTSEstimator
 from pts.model.predictor import Predictor
 from pts.model.forecast import Forecast
+from pts.model.quantile import Quantile

pts/model/deepar/__init__.py

@@ -0,0 +1 @@
+from .deepar_estimator import DeepAREstimator

pts/model/deepar/deepar_estimator.py

@@ -0,0 +1,27 @@
+from pts.model import PTSEstimator
+
+class DeepAREstimator(PTSEstimator):
+    def __init__(self,
+        freq: str,
+        prediction_length: 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_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(),
+        scaling: bool = True,
+        lags_seq: Optional[List[int]] = None,
+        time_features: Optional[List[TimeFeature]] = None,
+        num_parallel_samples: int = 100,
+    ) -> None:
+        super().__init__(trainer=trainer)
+        
+        
+

pts/model/deepar/deepar_network.py

@@ -0,0 +1,5 @@
+import torch
+import torch.nn as nn
+
+class DeepARNetwork(nn.Module):
+    pass

pts/model/estimator.py

@@ -10,7 +10,7 @@ import torch
 import torch.nn as nn
 
 from .predictor import Predictor
-
+from .utils import get_module_forward_input_names
 
 class Estimator(ABC):
     prediction_length: int
@@ -48,9 +48,9 @@ class TrainOutput(NamedTuple):
 
 class PTSEstimator(Estimator):
     def __init__(self, trainer: Trainer,
-                 float_type: np.dtype = np.float32) -> None:
+                 dtype: np.dtype = np.float32) -> None:
         self.trainer = trainer
-        self.float_type = float_type
+        self.dtype = dtype
         
 
     @abstractmethod
@@ -103,7 +103,7 @@ class PTSEstimator(Estimator):
             batch_size=self.trainer.batch_size,
             num_batches_per_epoch=self.trainer.num_batches_per_epoch,
             device=self.trainer.device,
-            float_type=self.float_type,
+            dtype=self.dtype,
         )
 
         # ensure that the training network is created on the same device
@@ -111,18 +111,15 @@ class PTSEstimator(Estimator):
 
         self.trainer(
             net=trained_net,
-            input_names=get_hybrid_forward_input_names(trained_net),
+            input_names=get_module_forward_input_names(trained_net),
             train_iter=training_data_loader,
         )
 
-        with self.trainer.ctx:
-            # ensure that the prediction network is created within the same MXNet
-            # context as the one that was used during training
-            return TrainOutput(
-                transformation=transformation,
-                trained_net=trained_net,
-                predictor=self.create_predictor(transformation, trained_net),
-            )
+        return TrainOutput(
+            transformation=transformation,
+            trained_net=trained_net,
+            predictor=self.create_predictor(transformation, trained_net),
+        )
 
     def train(self, training_data: Dataset) -> Predictor:
         return self.train_model(training_data).predictor

pts/model/forecast.py

@@ -1,2 +1,395 @@
-class Forecast():
-    pass
+from abc import ABC, abstractmethod
+from typing import Optional, Dict, Enum, Set, List
+
+import numpy as np
+import pandas as pd
+
+import torch
+from torch.distributions.distribution import Distributions
+
+from .quantile import Quantile
+
+
+class OutputType(str, Enum):
+    mean = "mean"
+    samples = "samples"
+    quantiles = "quantiles"
+
+
+class Config():
+    output_types: Set[OutputType] = {"quantiles", "mean"}
+    quantiles: List[str] = ["0.1", "0.5", "0.9"]
+
+
+class Forecast(ABC):
+    start_date: pd.Timestamp
+    freq: str
+    item_id: Optional[str]
+    info: Optional[Dict]
+    prediction_length: int
+    mean: np.ndarray
+    _index = None
+
+    @abstractmethod
+    def quantile(self, q: Union[float, str]) -> np.ndarray:
+        """
+        Computes a quantile from the predicted distribution.
+
+        Parameters
+        ----------
+        q
+            Quantile to compute.
+
+        Returns
+        -------
+        numpy.ndarray
+            Value of the quantile across the prediction range.
+        """
+        pass
+
+    @abstractmethod
+    def dim(self) -> int:
+        """
+        Returns the dimensionality of the forecast object.
+        """
+        pass
+
+    @abstractmethod
+    def copy_dim(self, dim: int):
+        """
+        Returns a new Forecast object with only the selected sub-dimension.
+
+        Parameters
+        ----------
+        dim
+            The returned forecast object will only represent this dimension.
+        """
+        pass
+
+    def as_json_dict(self, config: "Config") -> dict:
+        result = {}
+
+        if OutputType.mean in config.output_types:
+            result["mean"] = self.mean.tolist()
+
+        if OutputType.quantiles in config.output_types:
+            quantiles = map(Quantile.parse, config.quantiles)
+
+            result["quantiles"] = {
+                quantile.name: self.quantile(quantile.value).tolist()
+                for quantile in quantiles
+            }
+
+        if OutputType.samples in config.output_types:
+            result["samples"] = []
+
+        return result
+
+
+class SampleForecast(Forecast):
+    """
+    A `Forecast` object, where the predicted distribution is represented
+    internally as samples.
+
+    Parameters
+    ----------
+    samples
+        Array of size (num_samples, prediction_length)
+    start_date
+        start of the forecast
+    freq
+        forecast frequency
+    info
+        additional information that the forecaster may provide e.g. estimated
+        parameters, number of iterations ran etc.
+    """
+    @validated()
+    def __init__(
+            self,
+            samples: Union[torch.Tensor, np.ndarray],
+            start_date,
+            freq,
+            item_id: Optional[str] = None,
+            info: Optional[Dict] = None,
+    ):
+        assert isinstance(
+            samples, (np.ndarray, torch.Tensor
+                      )), "samples should be either a numpy array or an torch tensor"
+        assert (
+            len(np.shape(samples)) == 2 or len(np.shape(samples)) == 3
+        ), "samples should be a 2-dimensional or 3-dimensional array. Dimensions found: {}".format(
+            len(np.shape(samples)))
+        self.samples = (samples if (isinstance(samples, np.ndarray)) else
+                        samples.numpy())
+        self._sorted_samples_value = None
+        self._mean = None
+        self._dim = None
+        self.item_id = item_id
+        self.info = info
+
+        assert isinstance(
+            start_date,
+            pd.Timestamp), "start_date should be a pandas Timestamp object"
+        self.start_date = start_date
+
+        assert isinstance(freq, str), "freq should be a string"
+        self.freq = freq
+
+    @property
+    def _sorted_samples(self):
+        if self._sorted_samples_value is None:
+            self._sorted_samples_value = np.sort(self.samples, axis=0)
+        return self._sorted_samples_value
+
+    @property
+    def num_samples(self):
+        """
+        The number of samples representing the forecast.
+        """
+        return self.samples.shape[0]
+
+    @property
+    def prediction_length(self):
+        """
+        Time length of the forecast.
+        """
+        return self.samples.shape[-1]
+
+    @property
+    def mean(self):
+        """
+        Forecast mean.
+        """
+        if self._mean is not None:
+            return self._mean
+        else:
+            return np.mean(self.samples, axis=0)
+
+    @property
+    def mean_ts(self):
+        """
+        Forecast mean, as a pandas.Series object.
+        """
+        return pd.Series(self.index, self.mean)
+
+    def quantile(self, q):
+        q = Quantile.parse(q).value
+        sample_idx = int(np.round((self.num_samples - 1) * q))
+        return self._sorted_samples[sample_idx, :]
+
+    def copy_dim(self, dim: int):
+        if len(self.samples.shape) == 2:
+            samples = self.samples
+        else:
+            target_dim = self.samples.shape[2]
+            assert dim < target_dim, (
+                f"must set 0 <= dim < target_dim, but got dim={dim},"
+                f" target_dim={target_dim}")
+            samples = self.samples[:, :, dim]
+
+        return SampleForecast(samples=samples,
+                              start_date=self.start_date,
+                              freq=self.freq,
+                              item_id=self.item_id,
+                              info=self.info)
+
+    def dim(self) -> int:
+        if self._dim is not None:
+            return self._dim
+        else:
+            if len(self.samples.shape) == 2:
+                # univariate target
+                # shape: (num_samples, prediction_length)
+                return 1
+            else:
+                # multivariate target
+                # shape: (num_samples, prediction_length, target_dim)
+                return self.samples.shape[2]
+
+    def as_json_dict(self, config: "Config") -> dict:
+        result = super().as_json_dict(config)
+
+        if OutputType.samples in config.output_types:
+            result["samples"] = self.samples.tolist()
+
+        return result
+
+    def __repr__(self):
+        return ", ".join([
+            f"SampleForecast({self.samples!r})",
+            f"{self.start_date!r}",
+            f"{self.freq!r}",
+            f"item_id={self.item_id!r}",
+            f"info={self.info!r})",
+        ])
+
+
+class QuantileForecast(Forecast):
+    """
+    A Forecast that contains arrays (i.e. time series) for quantiles and mean
+
+    Parameters
+    ----------
+    forecast_arrays
+        An array of forecasts
+    start_date
+        start of the forecast
+    freq
+        forecast frequency
+    forecast_keys
+        A list of quantiles of the form '0.1', '0.9', etc.,
+        and potentially 'mean'. Each entry corresponds to one array in
+        forecast_arrays.
+    info
+        additional information that the forecaster may provide e.g. estimated
+        parameters, number of iterations ran etc.
+    """
+    def __init__(
+            self,
+            forecast_arrays: np.ndarray,
+            start_date: pd.Timestamp,
+            freq: str,
+            forecast_keys: List[str],
+            item_id: Optional[str] = None,
+            info: Optional[Dict] = None,
+    ):
+        self.forecast_array = forecast_arrays
+        self.start_date = pd.Timestamp(start_date, freq=freq)
+        self.freq = freq
+
+        # normalize keys
+        self.forecast_keys = [
+            Quantile.from_str(key).name if key != "mean" else key
+            for key in forecast_keys
+        ]
+        self.item_id = item_id
+        self.info = info
+        self._dim = None
+
+        shape = self.forecast_array.shape
+        assert shape[0] == len(self.forecast_keys), (
+            f"The forecast_array (shape={shape} should have the same "
+            f"length as the forecast_keys (len={len(self.forecast_keys)}).")
+        self.prediction_length = shape[-1]
+        self._forecast_dict = {
+            k: self.forecast_array[i]
+            for i, k in enumerate(self.forecast_keys)
+        }
+
+        self._nan_out = np.array([np.nan] * self.prediction_length)
+
+    def quantile(self, q: Union[float, str]) -> np.ndarray:
+        q_str = Quantile.parse(q).name
+        # We return nan here such that evaluation runs through
+        return self._forecast_dict.get(q_str, self._nan_out)
+
+    @property
+    def mean(self):
+        """
+        Forecast mean.
+        """
+        return self._forecast_dict.get("mean", self._nan_out)
+
+    def dim(self) -> int:
+        if self._dim is not None:
+            return self._dim
+        else:
+            if (len(self.forecast_array.shape) == 2
+                ):  # 1D target. shape: (num_samples, prediction_length)
+                return 1
+            else:
+                return self.forecast_array.shape[
+                    1]  # 2D target. shape: (num_samples, target_dim, prediction_length)
+
+    def __repr__(self):
+        return ", ".join([
+            f"QuantileForecast({self.forecast_array!r})",
+            f"start_date={self.start_date!r}",
+            f"freq={self.freq!r}",
+            f"forecast_keys={self.forecast_keys!r}",
+            f"item_id={self.item_id!r}",
+            f"info={self.info!r})",
+        ])
+
+
+# class DistributionForecast(Forecast):
+#     """
+#     A `Forecast` object that uses a distribution directly.
+#     This can for instance be used to represent marginal probability
+#     distributions for each time point -- although joint distributions are
+#     also possible, e.g. when using MultiVariateGaussian).
+
+#     Parameters
+#     ----------
+#     distribution
+#         Distribution object. This should represent the entire prediction
+#         length, i.e., if we draw `num_samples` samples from the distribution,
+#         the sample shape should be
+
+#            samples = trans_dist.sample(num_samples)
+#            samples.shape -> (num_samples, prediction_length)
+
+#     start_date
+#         start of the forecast
+#     freq
+#         forecast frequency
+#     info
+#         additional information that the forecaster may provide e.g. estimated
+#         parameters, number of iterations ran etc.
+#     """
+#     @validated()
+#     def __init__(
+#             self,
+#             distribution: Distribution,
+#             start_date,
+#             freq,
+#             item_id: Optional[str] = None,
+#             info: Optional[Dict] = None,
+#     ):
+#         self.distribution = distribution
+#         self.shape = (self.distribution.batch_shape +
+#                       self.distribution.event_shape)
+#         self.prediction_length = self.shape[0]
+#         self.item_id = item_id
+#         self.info = info
+
+#         assert isinstance(
+#             start_date,
+#             pd.Timestamp), "start_date should be a pandas Timestamp object"
+#         self.start_date = start_date
+
+#         assert isinstance(freq, str), "freq should be a string"
+#         self.freq = freq
+#         self._mean = None
+
+#     @property
+#     def mean(self):
+#         """
+#         Forecast mean.
+#         """
+#         if self._mean is not None:
+#             return self._mean
+#         else:
+#             self._mean = self.distribution.mean.asnumpy()
+#             return self._mean
+
+#     @property
+#     def mean_ts(self):
+#         """
+#         Forecast mean, as a pandas.Series object.
+#         """
+#         return pd.Series(self.index, self.mean)
+
+#     def quantile(self, level):
+#         level = Quantile.parse(level).value
+#         q = self.distribution.quantile(mx.nd.array([level])).asnumpy()[0]
+#         return q
+
+#     def to_sample_forecast(self, num_samples: int = 200) -> SampleForecast:
+#         return SampleForecast(
+#             samples=self.distribution.sample(num_samples),
+#             start_date=self.start_date,
+#             freq=self.freq,
+#             item_id=self.item_id,
+#             info=self.info,
+#         )

pts/model/quantile.py

@@ -0,0 +1,84 @@
+from typing import NamedTuple, Union
+import re
+
+class Quantile(NamedTuple):
+    value: float
+    name: str
+
+    @property
+    def loss_name(self):
+        return f"QuantileLoss[{self.name}]"
+
+    @property
+    def weighted_loss_name(self):
+        return f"wQuantileLoss[{self.name}]"
+
+    @property
+    def coverage_name(self):
+        return f"Coverage[{self.name}]"
+
+    @classmethod
+    def checked(cls, value: float, name: str) -> "Quantile":
+        if not 0 <= value <= 1:
+            raise Exception(
+                f"quantile value should be in [0, 1] but found {value}")
+
+        return Quantile(value, name)
+
+    @classmethod
+    def from_float(cls, quantile: float) -> "Quantile":
+        assert isinstance(quantile, float)
+        return cls.checked(value=quantile, name=str(quantile))
+
+    @classmethod
+    def from_str(cls, quantile: str) -> "Quantile":
+        assert isinstance(quantile, str)
+        try:
+            return cls.checked(value=float(quantile), name=quantile)
+        except ValueError:
+            m = re.match(r"^p(\d{2})$", quantile)
+
+            if m is None:
+                raise Exception(
+                    "Quantile string should be of the form "
+                    f'"p10", "p50", ... or "0.1", "0.5", ... but found {quantile}'
+                )
+            else:
+                quantile: float = int(m.group(1)) / 100
+                return cls(value=quantile, name=str(quantile))
+
+    @classmethod
+    def parse(cls, quantile: Union["Quantile", float, str]) -> "Quantile":
+        """Produces equivalent float and string representation of a given
+        quantile level.
+
+        >>> Quantile.parse(0.1)
+        Quantile(value=0.1, name='0.1')
+
+        >>> Quantile.parse('0.2')
+        Quantile(value=0.2, name='0.2')
+
+        >>> Quantile.parse('0.20')
+        Quantile(value=0.2, name='0.20')
+
+        >>> Quantile.parse('p99')
+        Quantile(value=0.99, name='0.99')
+
+        Parameters
+        ----------
+        quantile
+            Quantile, can be a float a str representing a float e.g. '0.1' or a
+            quantile string of the form 'p0.1'.
+
+        Returns
+        -------
+        Quantile
+            A tuple containing both a float and a string representation of the
+            input quantile level.
+        """
+        if isinstance(quantile, Quantile):
+            return quantile
+        elif isinstance(quantile, float):
+            return cls.from_float(quantile)
+        else:
+            return cls.from_str(quantile)

pts/model/utils.py

@@ -0,0 +1,7 @@
+import inspect
+
+import torch.nn as nn
+
+def get_module_forward_input_names(module: nn.Module):
+    params = inspect.signature(module.forward).parameters
+    return list(params)