Add FEDformer as an imputation model

environment-dev.yml

@@ -12,6 +12,7 @@ dependencies:
  - conda-forge::h5py
  - conda-forge::numpy
  - conda-forge::scipy
+ - conda-forge::sympy
  - conda-forge::python
  - conda-forge::einops
  - conda-forge::pandas
@@ -26,8 +27,8 @@ dependencies:
 
  # optional
  - pyg::pyg
- - pyg::pytorch-scatter
  - pyg::pytorch-sparse
+ - pyg::pytorch-scatter
 
  # test
  - conda-forge::pytest-cov

pypots/imputation/__init__.py

@@ -14,6 +14,7 @@
 from .transformer import Transformer
 from .timesnet import TimesNet
 from .etsformer import ETSformer
+from .fedformer import FEDformer
 from .crossformer import Crossformer
 from .autoformer import Autoformer
 from .dlinear import DLinear
@@ -30,6 +31,7 @@
  "SAITS",
  "Transformer",
  "ETSformer",
+ "FEDformer",
  "Crossformer",
  "TimesNet",
  "PatchTST",

pypots/imputation/fedformer/__init__.py

@@ -0,0 +1,17 @@
+"""
+The package of the partially-observed time-series imputation model FEDformer.
+
+Refer to the paper "Wu, H., Xu, J., Wang, J., & Long, M. (2021).
+FEDformer: 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 FEDformer
+
+__all__ = [
+ "FEDformer",
+]

pypots/imputation/fedformer/data.py

@@ -0,0 +1,24 @@
+"""
+Dataset class for FEDformer.
+"""
+
+# Created by Wenjie Du <wenjay.du@gmail.com>
+# License: BSD-3-Clause
+
+from typing import Union
+
+from ..saits.data import DatasetForSAITS
+
+
+class DatasetForFEDformer(DatasetForSAITS):
+ """Actually FEDformer 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)