Add PatchTST as an imputation model

pypots/imputation/__init__.py

@@ -14,6 +14,7 @@
 from .transformer import Transformer
 from .timesnet import TimesNet
 from .autoformer import Autoformer
+from .patchtst import PatchTST
 from .usgan import USGAN
 
 # naive imputation methods
@@ -26,6 +27,7 @@
     "SAITS",
     "Transformer",
     "TimesNet",
+    "PatchTST",
     "Autoformer",
     "BRITS",
     "MRNN",

pypots/imputation/autoformer/data.py

@@ -1,5 +1,5 @@
 """
-Dataset class for TimesNet.
+Dataset class for Autoformer.
 """
 
 # Created by Wenjie Du <wenjay.du@gmail.com>

pypots/imputation/autoformer/model.py

@@ -1,5 +1,5 @@
 """
-The implementation of Transformer for the partially-observed time-series imputation task.
+The implementation of Autoformer for the partially-observed time-series imputation task.
 
 Refer to the paper "Wu, H., Xu, J., Wang, J., & Long, M. (2021).
 Autoformer: Decomposition transformers with auto-correlation for long-term series forecasting. NeurIPS 2021.".
@@ -31,7 +31,7 @@
 
 class Autoformer(BaseNNImputer):
     """The PyTorch implementation of the Autoformer model.
-    TimesNet is originally proposed by Wu et al. in :cite:`wu2021autoformer`.
+    Autoformer is originally proposed by Wu et al. in :cite:`wu2021autoformer`.
 
     Parameters
     ----------
@@ -56,7 +56,7 @@ class Autoformer(BaseNNImputer):
     factor :
         The factor of the auto correlation mechanism for the Autoformer model.
 
-    moving_avg_kernel_size :
+    moving_avg_window_size :
         The window size of moving average.
 
     dropout :
@@ -120,7 +120,7 @@ def __init__(
         d_model: int,
         d_ffn: int,
         factor: int,
-        moving_avg_kernel_size: int,
+        moving_avg_window_size: int,
         dropout: float = 0,
         batch_size: int = 32,
         epochs: int = 100,
@@ -149,7 +149,7 @@ def __init__(
         self.d_model = d_model
         self.d_ffn = d_ffn
         self.factor = factor
-        self.moving_avg_kernel_size = moving_avg_kernel_size
+        self.moving_avg_window_size = moving_avg_window_size
         self.dropout = dropout
 
         # set up the model
@@ -161,7 +161,7 @@ def __init__(
             self.d_model,
             self.d_ffn,
             self.factor,
-            self.moving_avg_kernel_size,
+            self.moving_avg_window_size,
             self.dropout,
         )
         self._send_model_to_given_device()

pypots/imputation/autoformer/modules/core.py

@@ -29,7 +29,7 @@ def __init__(
         d_model,
         d_ffn,
         factor,
-        moving_avg_kernel_size,
+        moving_avg_window_size,
         dropout,
         activation="relu",
         output_attention=False,
@@ -38,7 +38,7 @@ def __init__(
 
         self.seq_len = n_steps
         self.n_layers = n_layers
-        self.series_decomp = SeriesDecompositionBlock(moving_avg_kernel_size)
+        self.series_decomp = SeriesDecompositionBlock(moving_avg_window_size)
         self.enc_embedding = DataEmbedding_wo_Pos(
             n_features,
             d_model,
@@ -54,7 +54,7 @@ def __init__(
                     ),
                     d_model,
                     d_ffn,
-                    moving_avg_kernel_size,
+                    moving_avg_window_size,
                     dropout,
                     activation,
                 )

pypots/imputation/patchtst/__init__.py

@@ -0,0 +1,17 @@
+"""
+The package of the partially-observed time-series imputation model PatchTST.
+
+Refer to the paper "Wu, H., Xu, J., Wang, J., & Long, M. (2021).
+PatchTST: Decomposition transformers with auto-correlation for long-term series forecasting. NeurIPS 2021.".
+
+"""
+
+# Created by Wenjie Du <wenjay.du@gmail.com>
+# License: BSD-3-Clause
+
+
+from .model import PatchTST
+
+__all__ = [
+    "PatchTST",
+]

pypots/imputation/patchtst/data.py

@@ -0,0 +1,24 @@
+"""
+Dataset class for PatchTST.
+"""
+
+# Created by Wenjie Du <wenjay.du@gmail.com>
+# License: BSD-3-Clause
+
+from typing import Union
+
+from ..saits.data import DatasetForSAITS
+
+
+class DatasetForPatchTST(DatasetForSAITS):
+    """Actually PatchTST uses the same data strategy as SAITS, needs MIT for training."""
+
+    def __init__(
+        self,
+        data: Union[dict, str],
+        return_X_ori: bool,
+        return_labels: bool,
+        file_type: str = "h5py",
+        rate: float = 0.2,
+    ):
+        super().__init__(data, return_X_ori, return_labels, file_type, rate)