Merge pull request #507 from WenjieDu/dev

Add TEFN and implement it as an imputation model

docs/index.rst

@@ -133,6 +133,8 @@ The paper references are all listed at the bottom of this readme file.
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Type           | Algorithm                                                 | IMPU | FORE | CLAS | CLUS | ANOD | Year - Venue          |
 +================+===========================================================+======+======+======+======+======+=======================+
+| Neural Net     | TEFN🧑‍🔧           :cite:`zhan2024tefn`                  |  ✅  |      |      |      |      | ``2024 - arXiv``      |
++----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Neural Net     | TimeMixer      :cite:`wang2024timemixer`                  |  ✅  |      |      |      |      | ``2024 - ICLR``       |
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Neural Net     | iTransformer🧑‍🔧 :cite:`liu2024itransformer`             |  ✅  |      |      |      |      | ``2024 - ICLR``       |

docs/references.bib

@@ -763,3 +763,10 @@ @article{bai2018tcn
 journal={arXiv preprint arXiv:1803.01271},
 year={2018}
 }
+
+@article{zhan2024tefn,
+  title={Time Evidence Fusion Network: Multi-source View in Long-Term Time Series Forecasting},
+  author={Zhan, Tianxiang and He, Yuanpeng and Li, Zhen and Deng, Yong},
+  journal={arXiv preprint arXiv:2405.06419},
+  year={2024}
+}

pypots/imputation/__init__.py

@@ -5,7 +5,6 @@
 # Created by Wenjie Du <[email protected]>
 # License: BSD-3-Clause
 
-# neural network imputation methods
 from .brits import BRITS
 from .csdi import CSDI
 from .gpvae import GPVAE
@@ -44,6 +43,7 @@
 from .mean import Mean
 from .median import Median
 from .lerp import Lerp
+from .tefn import TEFN
 
 __all__ = [
     # neural network imputation methods
@@ -84,4 +84,5 @@
     "Mean",
     "Median",
     "Lerp",
+    "TEFN",
 ]

pypots/imputation/tefn/__init__.py

@@ -0,0 +1,24 @@
+"""
+The package of the forecasting model TEFN.
+
+Refer to the paper
+`Tianxiang Zhan, Yuanpeng He, Yong Deng, and Zhen Li.
+Time Evidence Fusion Network: Multi-source View in Long-Term Time Series Forecasting.
+In Arxiv, 2024.
+<https://arxiv.org/abs/2405.06419>`_
+
+Notes
+-----
+This implementation is transfered from the official one https://github.com/ztxtech/Time-Evidence-Fusion-Network
+
+"""
+
+# Created by Tianxiang Zhan <[email protected]>
+# License: BSD-3-Clause
+
+
+from .model import TEFN
+
+__all__ = [
+    "TEFN",
+]

pypots/imputation/tefn/core.py

@@ -0,0 +1,59 @@
+"""
+
+"""
+
+# Created by Tianxiang Zhan <[email protected]>
+# License: BSD-3-Clause
+
+import torch.nn as nn
+
+from ...nn.functional import nonstationary_norm, nonstationary_denorm
+from ...nn.modules.tefn import BackboneTEFN
+from ...utils.metrics import calc_mse
+
+
+class _TEFN(nn.Module):
+    def __init__(
+        self,
+        n_steps,
+        n_features,
+        n_fod,
+        apply_nonstationary_norm,
+    ):
+        super().__init__()
+
+        self.seq_len = n_steps
+        self.n_fod = n_fod
+        self.apply_nonstationary_norm = apply_nonstationary_norm
+
+        self.model = BackboneTEFN(
+            n_steps,
+            n_features,
+            n_fod,
+        )
+
+    def forward(self, inputs: dict, training: bool = True) -> dict:
+        X, missing_mask = inputs["X"], inputs["missing_mask"]
+
+        if self.apply_nonstationary_norm:
+            # Normalization from Non-stationary Transformer
+            X, means, stdev = nonstationary_norm(X, missing_mask)
+
+        # TEFN processing
+        out = self.model(X)
+
+        if self.apply_nonstationary_norm:
+            # De-Normalization from Non-stationary Transformer
+            out = nonstationary_denorm(out, means, stdev)
+
+        imputed_data = missing_mask * X + (1 - missing_mask) * out
+        results = {
+            "imputed_data": imputed_data,
+        }
+
+        if training:
+            # `loss` is always the item for backward propagating to update the model
+            loss = calc_mse(out, inputs["X_ori"], inputs["indicating_mask"])
+            results["loss"] = loss
+
+        return results

pypots/imputation/tefn/data.py

@@ -0,0 +1,24 @@
+"""
+Dataset class for the imputation model TEFN.
+"""
+
+# Created by Tianxiang Zhan <[email protected]>
+# License: BSD-3-Clause
+
+from typing import Union
+
+from ..saits.data import DatasetForSAITS
+
+
+class DatasetForTEFN(DatasetForSAITS):
+    """Actually TEFN uses the same data strategy as SAITS, needs MIT for training."""
+
+    def __init__(
+        self,
+        data: Union[dict, str],
+        return_X_ori: bool,
+        return_y: bool,
+        file_type: str = "hdf5",
+        rate: float = 0.2,
+    ):
+        super().__init__(data, return_X_ori, return_y, file_type, rate)