Add FITS imputation model and update docs

README.md

@@ -121,6 +121,7 @@ The paper references and links are all listed at the bottom of this file.
 |:--------------|:---------------------------------------------------------------------------------------------------------------------------------|:--------:|:--------:|:--------:|:--------:|:--------:|:-------------------|
 | LLM           | <a href="https://time-series.ai"><img src="https://time-series.ai/static/figs/robot.svg" width="26px"> Time-Series.AI</a>  [^36] |    ✅     |    ✅     |    ✅     |    ✅     |    ✅     | `Later in 2024`    |
 | Neural Net    | TEFN🧑‍🔧[^39]                                                                                                                   |    ✅     |          |          |          |          | `2024 - arXiv`     |
+| Neural Net    | FITS🧑‍🔧[^41]                                                                                                                   |    ✅     |          |          |          |          | `2024 - ICLR`      |
 | Neural Net    | TimeMixer[^37]                                                                                                                   |    ✅     |          |          |          |          | `2024 - ICLR`      |
 | Neural Net    | iTransformer🧑‍🔧[^24]                                                                                                           |    ✅     |          |          |          |          | `2024 - ICLR`      |
 | Neural Net    | ModernTCN[^38]                                                                                                                   |    ✅     |          |          |          |          | `2024 - ICLR`      |
@@ -501,7 +502,10 @@ Time-Series.AI</a>
 [^38]: Luo, D., & Wang X. (2024).
 [ModernTCN: A Modern Pure Convolution Structure for General Time Series Analysis](https://openreview.net/forum?id=vpJMJerXHU).
 *ICLR 2024*.
-[^39]: Zhan, T., He, Y., Li, Z., & Deng, Y. (2024).
+[^39]: Zhan, T., He, Y., Deng, Y., Li, Z., Du, W., & Wen, Q. (2024).
 [Time Evidence Fusion Network: Multi-source View in Long-Term Time Series Forecasting](https://arxiv.org/abs/2405.06419).
 *arXiv 2024*.
-[^40]: [Wikipedia: Linear interpolation](https://en.wikipedia.org/wiki/Linear_interpolation)
+[^40]: [Wikipedia: Linear interpolation](https://en.wikipedia.org/wiki/Linear_interpolation)
+[^41]: Xu, Z., Zeng, A., & Xu, Q. (2024).
+[FITS: Modeling Time Series with 10k parameters](https://openreview.net/forum?id=bWcnvZ3qMb).
+*ICLR 2024*.

README_zh.md

@@ -106,6 +106,7 @@ PyPOTS当前支持多变量POTS数据的插补, 预测, 分类, 聚类以及异
 |:--------------|:---------------------------------------------------------------------------------------------------------------------------------|:------:|:------:|:------:|:------:|:--------:|:-------------------|
 | LLM           | <a href="https://time-series.ai"><img src="https://time-series.ai/static/figs/robot.svg" width="26px"> Time-Series.AI</a>  [^36] |   ✅    |   ✅    |   ✅    |   ✅    |    ✅     | `Later in 2024`    |
 | Neural Net    | TEFN🧑‍🔧[^39]                                                                                                                   |   ✅    |        |        |        |          | `2024 - arXiv`     |
+| Neural Net    | FITS🧑‍🔧[^41]                                                                                                                   |   ✅    |        |        |        |          | `2024 - ICLR`      |
 | Neural Net    | TimeMixer[^37]                                                                                                                   |   ✅    |        |        |        |          | `2024 - ICLR`      |
 | Neural Net    | iTransformer🧑‍🔧[^24]                                                                                                           |   ✅    |        |        |        |          | `2024 - ICLR`      |
 | Neural Net    | ModernTCN[^38]                                                                                                                   |   ✅    |        |        |        |          | `2024 - ICLR`      |
@@ -474,7 +475,10 @@ Time-Series.AI</a>
 [^38]: Luo, D., & Wang X. (2024).
 [ModernTCN: A Modern Pure Convolution Structure for General Time Series Analysis](https://openreview.net/forum?id=vpJMJerXHU).
 *ICLR 2024*.
-[^39]: Zhan, T., He, Y., Li, Z., & Deng, Y. (2024).
+[^39]: Zhan, T., He, Y., Deng, Y., Li, Z., Du, W., & Wen, Q. (2024).
 [Time Evidence Fusion Network: Multi-source View in Long-Term Time Series Forecasting](https://arxiv.org/abs/2405.06419).
 *arXiv 2024*.
-[^40]: [Wikipedia: Linear interpolation](https://en.wikipedia.org/wiki/Linear_interpolation)
+[^40]: [Wikipedia: Linear interpolation](https://en.wikipedia.org/wiki/Linear_interpolation)
+[^41]: Xu, Z., Zeng, A., & Xu, Q. (2024).
+[FITS: Modeling Time Series with 10k parameters](https://openreview.net/forum?id=bWcnvZ3qMb).
+*ICLR 2024*.

docs/conf.py

@@ -62,7 +62,7 @@
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3", None),
     "sphinx": ("https://www.sphinx-doc.org/en/master", None),
-    "torch": ("https://pytorch.org/docs/master/", None),
+    "torch": ("https://pytorch.org/docs/main/", None),
     "numpy": ("https://numpy.org/doc/stable/", None),
     "pandas": ("https://pandas.pydata.org/docs/", None),
 }

docs/index.rst

@@ -135,6 +135,8 @@ The paper references are all listed at the bottom of this readme file.
 +================+===========================================================+======+======+======+======+======+=======================+
 | Neural Net     | TEFN🧑‍🔧           :cite:`zhan2024tefn`                  |  ✅  |      |      |      |      | ``2024 - arXiv``      |
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
+| Neural Net     | FITS🧑‍🔧           :cite:`xu2024fits`                    |  ✅  |      |      |      |      | ``2024 - ICLR``       |
++----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Neural Net     | TimeMixer      :cite:`wang2024timemixer`                  |  ✅  |      |      |      |      | ``2024 - ICLR``       |
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Neural Net     | iTransformer🧑‍🔧 :cite:`liu2024itransformer`             |  ✅  |      |      |      |      | ``2024 - ICLR``       |

docs/pypots.imputation.rst

@@ -19,6 +19,24 @@ pypots.imputation.transformer
    :show-inheritance:
    :inherited-members:
 
+pypots.imputation.tefn
+------------------------------------
+
+.. automodule:: pypots.imputation.tefn
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :inherited-members:
+
+pypots.imputation.fits
+------------------------------------
+
+.. automodule:: pypots.imputation.fits
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :inherited-members:
+
 pypots.imputation.timemixer
 ------------------------------------
 

docs/references.bib

@@ -766,7 +766,7 @@ @article{bai2018tcn
 
 @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},
+author={Zhan, Tianxiang and He, Yuanpeng and Deng, Yong and Li, Zhen and Du, Wenjie and Wen, Qingsong},
 journal={arXiv preprint arXiv:2405.06419},
 year={2024}
 }

pypots/imputation/fits/__init__.py

@@ -0,0 +1,24 @@
+"""
+The package including the modules of FITS.
+
+Refer to the paper
+`Zhijian Xu, Ailing Zeng, and Qiang Xu.
+FITS: Modeling Time Series with 10k parameters.
+In The Twelfth International Conference on Learning Representations, 2024.
+<https://openreview.net/pdf?id=bWcnvZ3qMb>`_
+
+Notes
+-----
+This implementation is inspired by the official one https://github.com/VEWOXIC/FITS
+
+"""
+
+# Created by Wenjie Du <wenjay.du@gmail.com>
+# License: BSD-3-Clause
+
+
+from .model import FITS
+
+__all__ = [
+    "FITS",
+]

pypots/imputation/fits/core.py

@@ -0,0 +1,86 @@
+"""
+The core wrapper assembles the submodules of FITS imputation model
+and takes over the forward progress of the algorithm.
+"""
+
+# Created by Wenjie Du <wenjay.du@gmail.com>
+# License: BSD-3-Clause
+
+import torch.nn as nn
+
+from ...nn.functional import nonstationary_norm, nonstationary_denorm
+from ...nn.modules.fits import BackboneFITS
+from ...nn.modules.saits import SaitsLoss, SaitsEmbedding
+
+
+class _FITS(nn.Module):
+    def __init__(
+        self,
+        n_steps: int,
+        n_features: int,
+        cut_freq: int,
+        individual: bool,
+        ORT_weight: float = 1,
+        MIT_weight: float = 1,
+        apply_nonstationary_norm: bool = False,
+    ):
+        super().__init__()
+
+        self.n_steps = n_steps
+        self.apply_nonstationary_norm = apply_nonstationary_norm
+
+        self.saits_embedding = SaitsEmbedding(
+            n_features * 2,
+            n_features,
+            with_pos=False,
+        )
+        self.backbone = BackboneFITS(
+            n_steps,
+            n_features,
+            0,  # n_pred_steps is not used in the imputation task
+            cut_freq,
+            individual,
+        )
+
+        # for the imputation task, the output dim is the same as input dim
+        self.output_projection = nn.Linear(n_features, n_features)
+        self.saits_loss_func = SaitsLoss(ORT_weight, MIT_weight)
+
+    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)
+
+        # WDU: the original FITS paper isn't proposed for imputation task. Hence the model doesn't take
+        # the missing mask into account, which means, in the process, the model doesn't know which part of
+        # the input data is missing, and this may hurt the model's imputation performance. Therefore, I apply the
+        # SAITS embedding method to project the concatenation of features and masks into a hidden space, as well as
+        # the output layers to project back from the hidden space to the original space.
+        enc_out = self.saits_embedding(X, missing_mask)
+
+        # FITS encoder processing
+        enc_out = self.backbone(enc_out)
+        if self.apply_nonstationary_norm:
+            # De-Normalization from Non-stationary Transformer
+            enc_out = nonstationary_denorm(enc_out, means, stdev)
+
+        # project back the original data space
+        reconstruction = self.output_projection(enc_out)
+
+        imputed_data = missing_mask * X + (1 - missing_mask) * reconstruction
+        results = {
+            "imputed_data": imputed_data,
+        }
+
+        # if in training mode, return results with losses
+        if training:
+            X_ori, indicating_mask = inputs["X_ori"], inputs["indicating_mask"]
+            loss, ORT_loss, MIT_loss = self.saits_loss_func(reconstruction, X_ori, missing_mask, indicating_mask)
+            results["ORT_loss"] = ORT_loss
+            results["MIT_loss"] = MIT_loss
+            # `loss` is always the item for backward propagating to update the model
+            results["loss"] = loss
+
+        return results

pypots/imputation/fits/data.py

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