Merge pull request #539 from WenjieDu/dev

Implement SegRNN as an imputation model
WenjieDu · Oct 25, 2024 · 4a77213 · 4a77213
2 parents a4f1a72 + bf045b4
commit 4a77213
Show file tree

Hide file tree

Showing 23 changed files with 933 additions and 390 deletions.
diff --git a/README.md b/README.md
@@ -136,6 +136,8 @@ The paper references and links are all listed at the bottom of this file.
 | Neural Net    | MICN🧑‍🔧[^27]                                                                                                                   |    ✅     |          |          |          |          | `2023 - ICLR`      |
 | Neural Net    | DLinear🧑‍🔧[^17]                                                                                                                |    ✅     |          |          |          |          | `2023 - AAAI`      |
 | Neural Net    | TiDE🧑‍🔧[^28]                                                                                                                   |    ✅     |          |          |          |          | `2023 - TMLR`      |
+| Neural Net    | CSAI[^42]                                                                                                                        |    ✅     |          |          |          |          | `2023 - arXiv`     |
+| Neural Net    | SegRNN🧑‍🔧[^43]                                                                                                                 |    ✅     |          |          |          |          | `2023 - arXiv`     |
 | Neural Net    | SCINet🧑‍🔧[^30]                                                                                                                 |    ✅     |          |          |          |          | `2022 - NeurIPS`   |
 | Neural Net    | Nonstationary Tr.🧑‍🔧[^25]                                                                                                      |    ✅     |          |          |          |          | `2022 - NeurIPS`   |
 | Neural Net    | FiLM🧑‍🔧[^22]                                                                                                                   |    ✅     |          |          |          |          | `2022 - NeurIPS`   |
@@ -509,3 +511,9 @@ Time-Series.AI</a>
 [^41]: Xu, Z., Zeng, A., & Xu, Q. (2024).
 [FITS: Modeling Time Series with 10k parameters](https://openreview.net/forum?id=bWcnvZ3qMb).
 *ICLR 2024*.
+[^42]: Qian, L., Ibrahim, Z., Ellis, H. L., Zhang, A., Zhang, Y., Wang, T., & Dobson, R. (2023). 
+[Knowledge Enhanced Conditional Imputation for Healthcare Time-series](https://arxiv.org/abs/2312.16713).
+*arXiv 2023*.
+[^43]: Lin, S., Lin, W., Wu, W., Zhao, F., Mo, R., & Zhang, H. (2023). 
+[SegRNN: Segment Recurrent Neural Network for Long-Term Time Series Forecasting](https://arxiv.org/abs/2308.11200).
+*arXiv 2023*.
diff --git a/README_zh.md b/README_zh.md
@@ -121,6 +121,8 @@ PyPOTS当前支持多变量POTS数据的插补, 预测, 分类, 聚类以及异
 | Neural Net    | MICN🧑‍🔧[^27]                                                                                                                   |   ✅    |        |        |        |          | `2023 - ICLR`      |
 | Neural Net    | DLinear🧑‍🔧[^17]                                                                                                                |   ✅    |        |        |        |          | `2023 - AAAI`      |
 | Neural Net    | TiDE🧑‍🔧[^28]                                                                                                                   |   ✅    |        |        |        |          | `2023 - TMLR`      |
+| Neural Net    | CSAI[^42]                                                                                                                        |    ✅     |          |          |          |          | `2023 - arXiv`     |
+| Neural Net    | SegRNN🧑‍🔧[^43]                                                                                                                 |    ✅     |          |          |          |          | `2023 - arXiv`     |
 | Neural Net    | SCINet🧑‍🔧[^30]                                                                                                                 |   ✅    |        |        |        |          | `2022 - NeurIPS`   |
 | Neural Net    | Nonstationary Tr.🧑‍🔧[^25]                                                                                                      |   ✅    |        |        |        |          | `2022 - NeurIPS`   |
 | Neural Net    | FiLM🧑‍🔧[^22]                                                                                                                   |   ✅    |        |        |        |          | `2022 - NeurIPS`   |
@@ -482,3 +484,9 @@ Time-Series.AI</a>
 [^41]: Xu, Z., Zeng, A., & Xu, Q. (2024).
 [FITS: Modeling Time Series with 10k parameters](https://openreview.net/forum?id=bWcnvZ3qMb).
 *ICLR 2024*.
+[^42]: Qian, L., Ibrahim, Z., Ellis, H. L., Zhang, A., Zhang, Y., Wang, T., & Dobson, R. (2023). 
+[Knowledge Enhanced Conditional Imputation for Healthcare Time-series](https://arxiv.org/abs/2312.16713).
+*arXiv 2023*.
+[^43]: Lin, S., Lin, W., Wu, W., Zhao, F., Mo, R., & Zhang, H. (2023). 
+[SegRNN: Segment Recurrent Neural Network for Long-Term Time Series Forecasting](https://arxiv.org/abs/2308.11200).
+*arXiv 2023*.
diff --git a/docs/index.rst b/docs/index.rst
@@ -165,6 +165,10 @@ The paper references are all listed at the bottom of this readme file.
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Neural Net     | TiDE🧑‍🔧 :cite:`das2023tide`                             |  ✅  |      |      |      |      | ``2023 - TMLR``       |
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
+| Neural Net     | CSAI :cite:`qian2023csai`                                 |  ✅  |      |      |      |      | ``2023 - arXiv``      |
++----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
+| Neural Net     | SegRNN🧑‍🔧 :cite:`lin2023segrnn`                         |  ✅  |      |      |      |      | ``2023 - arXiv``      |
++----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Neural Net     | SCINet🧑‍🔧 :cite:`liu2022scinet`                         |  ✅  |      |      |      |      | ``2022 - NeurIPS``    |
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Neural Net     | Nonstationary Tr🧑‍🔧 :cite:`liu2022nonstationary`        |  ✅  |      |      |      |      | ``2022 - NeurIPS``    |

diff --git a/docs/pypots.imputation.rst b/docs/pypots.imputation.rst
@@ -28,6 +28,24 @@ pypots.imputation.tefn
    :show-inheritance:
    :inherited-members:
 
+pypots.imputation.csai
+------------------------------------
+
+.. automodule:: pypots.imputation.csai
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :inherited-members:
+
+pypots.imputation.segrnn
+------------------------------------
+
+.. automodule:: pypots.imputation.segrnn
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :inherited-members:
+
 pypots.imputation.fits
 ------------------------------------
 

diff --git a/docs/references.bib b/docs/references.bib
@@ -779,3 +779,9 @@ @inproceedings{jin2024timellm
 url={https://openreview.net/forum?id=Unb5CVPtae}
 }
 
+@article{qian2023csai,
+title={Knowledge Enhanced Conditional Imputation for Healthcare Time-series},
+author={Qian, Linglong and Ibrahim, Zina and Ellis, Hugh Logan and Zhang, Ao and Zhang, Yuezhou and Wang, Tao and Dobson, Richard},
+journal={arXiv preprint arXiv:2312.16713},
+year={2023}
+}
diff --git a/pypots/classification/csai/__init__.py b/pypots/classification/csai/__init__.py
@@ -17,4 +17,4 @@
 
 __all__ = [
     "CSAI",
-]
+]
diff --git a/pypots/classification/csai/core.py b/pypots/classification/csai/core.py
@@ -17,10 +17,10 @@
 #         self.bcelogits = nn.BCEWithLogitsLoss()
 
 #     def forward(self, y_score, y_out, targets, smooth=1):
-        
+
 #         #comment out if your model contains a sigmoid or equivalent activation layer
-#         # inputs = F.sigmoid(inputs)       
-        
+#         # inputs = F.sigmoid(inputs)
+
 #         #flatten label and prediction tensors
 #         BCE = self.bcelogits(y_out, targets)
 
@@ -30,23 +30,23 @@
 #         dice_loss = 1 - (2.*intersection + smooth)/(y_score.sum() + targets.sum() + smooth)
 
 #         Dice_BCE = BCE + dice_loss
-        
+
 #         return BCE, Dice_BCE
 
 
 class _BCSAI(nn.Module):
     def __init__(
-            self,
-            n_steps: int,
-            n_features: int,
-            rnn_hidden_size: int,
-            imputation_weight: float,
-            consistency_weight: float,
-            classification_weight: float,
-            n_classes: int,
-            step_channels: int,
-            dropout: float = 0.5,
-            intervals=None,
+        self,
+        n_steps: int,
+        n_features: int,
+        rnn_hidden_size: int,
+        imputation_weight: float,
+        consistency_weight: float,
+        classification_weight: float,
+        n_classes: int,
+        step_channels: int,
+        dropout: float = 0.5,
+        intervals=None,
     ):
         super().__init__()
         self.n_steps = n_steps
@@ -107,17 +107,17 @@ def forward(self, inputs: dict, training: bool = True) -> dict:
             b_classification_loss = F.nll_loss(torch.log(b_prediction), inputs["labels"])
             # f_classification_loss, _ = criterion(f_prediction, f_logits, inputs["labels"].unsqueeze(1).float())
             # b_classification_loss, _ = criterion(b_prediction, b_logits, inputs["labels"].unsqueeze(1).float())
-            classification_loss = (f_classification_loss + b_classification_loss) 
+            classification_loss = f_classification_loss + b_classification_loss
 
             loss = (
-                self.consistency_weight * consistency_loss + 
-                self.imputation_weight * reconstruction_loss +
-                self.classification_weight * classification_loss
+                self.consistency_weight * consistency_loss
+                + self.imputation_weight * reconstruction_loss
+                + self.classification_weight * classification_loss
             )
 
             results["loss"] = loss
             results["classification_loss"] = classification_loss
             results["f_reconstruction"] = f_reconstruction
             results["b_reconstruction"] = b_reconstruction
 
-        return results
+        return results
diff --git a/pypots/classification/csai/data.py b/pypots/classification/csai/data.py
@@ -6,22 +6,22 @@
 # License: BSD-3-Clause
 
 from typing import Union
-from ...imputation.csai.data import DatasetForCSAI as DatasetForCSAI_Imputation     
-
+from ...imputation.csai.data import DatasetForCSAI as DatasetForCSAI_Imputation
 
 
 class DatasetForCSAI(DatasetForCSAI_Imputation):
-    def __init__(self, 
-                 data: Union[dict, str],  
-                 file_type: str = "hdf5",
-                 return_y: bool = True,
-                 removal_percent: float = 0.0,
-                 increase_factor: float = 0.1,
-                 compute_intervals: bool = False,
-                 replacement_probabilities = None,
-                 normalise_mean : list = [],
-                 normalise_std: list = [],
-                 training: bool = True
+    def __init__(
+        self,
+        data: Union[dict, str],
+        file_type: str = "hdf5",
+        return_y: bool = True,
+        removal_percent: float = 0.0,
+        increase_factor: float = 0.1,
+        compute_intervals: bool = False,
+        replacement_probabilities=None,
+        normalise_mean: list = [],
+        normalise_std: list = [],
+        training: bool = True,
     ):
         super().__init__(
             data=data,
@@ -34,6 +34,5 @@ def __init__(self,
             replacement_probabilities=replacement_probabilities,
             normalise_mean=normalise_mean,
             normalise_std=normalise_std,
-            training=training
-       )
-
+            training=training,
+        )