Merge pull request #505 from ztxtech/main

Add TEFN model

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 <wenjay.du@gmail.com>
 # 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 <zhantianxianguestc@hotmail.com>
+# License: BSD-3-Clause
+
+
+from .model import TEFN
+
+__all__ = [
+ "TEFN",
+]

pypots/imputation/tefn/core.py

@@ -0,0 +1,59 @@
+"""
+
+"""
+
+# Created by Tianxiang Zhan <zhantianxianguestc@hotmail.com>
+# 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 <zhantianxianguestc@hotmail.com>
+# 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)