Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/architecture.pdf filter=lfs diff=lfs merge=lfs -text
+assets/spikf-go-architecture.png filter=lfs diff=lfs merge=lfs -text
+assets/supplementary.pdf filter=lfs diff=lfs merge=lfs -text

CITATION.cff

@@ -0,0 +1,12 @@
+cff-version: 1.2.0
+title: "SpikF-GO: Spiking Fourier Graph Operators for Multivariate Time Series Forecasting"
+message: "If you use this code, please cite our ECML PKDD 2026 paper."
+authors:
+  - family-names: Bakhshaliyev
+    given-names: Jafar
+  - family-names: Landwehr
+    given-names: Niels
+year: 2026
+conference: "ECML PKDD 2026"
+repository-code: "https://github.com/jafarbakhshaliyev/SpikF-GO"
+license: MIT

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Jafar Bakhshaliyev
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,194 @@
 ---
 license: mit
+tags:
+  - time-series
+  - forecasting
+  - spiking-neural-networks
+  - graph-neural-networks
+  - multivariate-time-series
 ---
+
+# SpikF-GO: Spiking Fourier Graph Operators for Multivariate Time Series Forecasting
+
+[![arXiv](https://img.shields.io/badge/arXiv-2606.13901-b31b1b.svg)](https://arxiv.org/abs/2606.13901)
+[![ECML PKDD 2026](https://img.shields.io/badge/ECML%20PKDD-2026-blue.svg)](https://arxiv.org/abs/2606.13901)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
+
+📄 **Paper (arXiv):** https://arxiv.org/abs/2606.13901  
+💻 **GitHub:** https://github.com/jafarbakhshaliyev/SpikF-GO
+
+Official implementation of **SpikF-GO: Spiking Fourier Graph Operators for Multivariate Time Series Forecasting**, accepted to the **ECML PKDD 2026 Research Track**.
+
+![SpikF-GO architecture](assets/spikf-go-architecture.png)
+
+---
+
+## Abstract
+
+SpikF-GO is a spiking neural architecture for multivariate time series forecasting. It combines the hypervariate graph formulation of FourierGNN with spike-driven Fourier-domain graph processing, enabling joint modeling of intra-series temporal dependencies, inter-series dependencies, and time-varying cross-variable interactions. The model introduces sparse frequency selection and Complex LIF-based spectral gating to preserve event-driven computation in the Fourier domain. We also provide **SpikF-GO w/ CPG**, which incorporates Central Pattern Generator-based positional signals for improved long-range temporal modeling.
+
+---
+
+## Key Contributions
+
+- **Graph-based SNN forecasting:** SpikF-GO brings hypervariate graph modeling into SNN-based multivariate time series forecasting.
+- **Spike-driven Fourier graph operators:** The model combines sparse frequency gating with Complex LIF-based spectral processing to preserve event-driven computation in the Fourier domain.
+- **Unified SNN benchmark:** We evaluate SpikF-GO against major SNN forecasting families under a common experimental protocol across eight benchmark datasets.
+- **Energy-aware forecasting:** SpikF-GO achieves competitive-to-superior forecasting performance while reducing theoretical energy consumption relative to FourierGNN.
+
+---
+
+## Related Library: SpikingTSF
+
+We also maintain **[SpikingTSF](https://github.com/spikora/SpikingTSF)**, a broader open-source benchmark library for spiking neural network-based time series forecasting. SpikingTSF unifies SNN forecasting architectures and ANN baselines under a common training and evaluation protocol across datasets, horizons, metrics, and random seeds.
+
+> **Note:** SpikingTSF is a benchmarking library and may not reproduce all experiments from this repository directly.
+
+---
+
+## Repository Structure
+
+```
+SpikF-GO/
+├── README.md
+├── LICENSE
+├── CITATION.cff
+├── requirements.txt
+├── train.py                  # main training & evaluation entry point
+├── model/                    # SpikF-GO + all baseline implementations
+├── utils/                    # shared utilities (metrics, helpers)
+├── data/
+│   └── data_loader.py        # dataset loading (raw files placed here at runtime)
+├── scripts/
+│   ├── ecg.sh
+│   ├── covid.sh
+│   ├── solar.sh
+│   ├── ecl.sh
+│   ├── traffic.sh
+│   ├── metr_la.sh
+│   ├── pems_bay.sh
+│   └── wiki.sh
+└── assets/
+    ├── spikf-go-architecture.png
+    └── supplementary.pdf
+```
+
+---
+
+## Environment Setup
+
+Create and activate a virtual environment:
+
+**Linux / macOS**
+```bash
+python3 -m venv venv
+source venv/bin/activate
+```
+
+**Windows**
+```bash
+python -m venv venv
+venv\Scripts\activate
+```
+
+Install dependencies:
+```bash
+pip install -r requirements.txt
+```
+
+Experiments were run with **PyTorch 2.5.1** on a single **NVIDIA RTX 4090**.
+
+---
+
+## Dataset
+
+Download the processed datasets from Figshare:
+
+https://figshare.com/s/7617530bce306584fe95?file=62576929
+
+Place all dataset files **directly** inside the `data/` folder (do **not** create subfolders):
+
+```
+SpikF-GO/
+├── data/
+│   ├── dataset_file_1
+│   ├── dataset_file_2
+│   └── ...
+├── model/
+├── scripts/
+└── train.py
+```
+
+---
+
+## Run Experiments
+
+Scripts are in `scripts/`, one per dataset:
+
+```bash
+bash scripts/ecg.sh
+bash scripts/covid.sh
+bash scripts/solar.sh
+bash scripts/ecl.sh
+bash scripts/traffic.sh
+bash scripts/metr_la.sh
+bash scripts/pems_bay.sh
+bash scripts/wiki.sh
+```
+
+Each script sets the exact hyperparameters used to produce the results reported in the paper.
+
+---
+
+## Supplementary Material
+
+Available at [`assets/supplementary.pdf`](assets/supplementary.pdf).
+
+---
+
+## Citation
+
+If you use this code or build on SpikF-GO, please cite our paper:
+
+**arXiv preprint:**
+```bibtex
+@misc{bakhshaliyev2026spikfgo,
+  title        = {SpikF-GO: Spiking Fourier Graph Operators for Multivariate Time Series Forecasting},
+  author       = {Bakhshaliyev, Jafar and Landwehr, Niels},
+  year         = {2026},
+  eprint       = {2606.13901},
+  archivePrefix= {arXiv},
+  primaryClass = {cs.LG},
+  url          = {https://arxiv.org/abs/2606.13901}
+}
+```
+
+**ECML PKDD 2026 proceedings:**
+```bibtex
+@inproceedings{bakhshaliyev2026spikfgo,
+  title     = {SpikF-GO: Spiking Fourier Graph Operators for Multivariate Time Series Forecasting},
+  author    = {Bakhshaliyev, Jafar and Landwehr, Niels},
+  booktitle = {ECML PKDD},
+  year      = {2026}
+}
+```
+
+See [`CITATION.cff`](CITATION.cff) for full citation metadata.
+
+---
+
+## Acknowledgements
+
+The baselines in `model/` build on prior work. We thank the authors for releasing their code; original licenses are respected.
+
+- **`SpikF.py`** — adapted from **SpikF** (Wu, Huo & Chen, *"SpikF: Spiking Fourier Network for Efficient Long-term Prediction"*, [ICML 2025 / PMLR v267](https://proceedings.mlr.press/v267/wu25m.html)).
+- **`TS_Former.py`, `TS_GRU.py`, `TS_TCN.py`** — adapted from **TS-LIF** (Feng et al., *"TS-LIF: A Temporal Segment Spiking Neuron Network for Time Series Forecasting"*, [arXiv:2503.05108](https://arxiv.org/abs/2503.05108)).
+- **`iSpikformer.py`, `SpikeGRU.py`** — adapted from **SeqSNN** (Lv et al., *"Efficient and Effective Time-Series Forecasting with Spiking Neural Networks"*, [arXiv:2402.01533](https://arxiv.org/abs/2402.01533)), [microsoft/SeqSNN](https://github.com/microsoft/SeqSNN).
+- **`SpikeRNN_CPG.py`, `SpikeTCN_CPG.py`, `Spikformer_CPG.py`** — CPG variants build on [arXiv:2405.14362](https://arxiv.org/abs/2405.14362) / [microsoft/SeqSNN](https://github.com/microsoft/SeqSNN).
+- **`FourierGNN.py`** — adapted from **FourierGNN**, [arXiv:2311.06190](https://arxiv.org/abs/2311.06190) / [aikunyi/FourierGNN](https://github.com/aikunyi/FourierGNN).
+
+---
+
+## License
+
+This project is released under the [MIT License](LICENSE).

assets/architecture.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b4851b633d30ef57fa79e8c1096ad660b31e8a9dad9376e7840b39ff3adc0a4
+size 194717

assets/supplementary.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:353d376afbe7e500c8af0a7cf21f927a95eeac9feb781a324991c5b2b063dc4e
+size 219947

data/data_loader.py

@@ -0,0 +1,243 @@
+import os
+import datetime
+import numpy as np
+import pandas as pd
+
+from torch.utils.data import Dataset
+from sklearn.preprocessing import StandardScaler
+
+
+
+def _split_with_overlap(data: np.ndarray, train_ratio: float, val_ratio: float, seq_len: int):
+    """
+    Time split with overlap for val/test to allow past context:
+      train: [0 : train_end)
+      val  : [train_end - seq_len : val_end)
+      test : [val_end - seq_len : T)
+    """
+    T = len(data)
+    train_end = int(T * train_ratio)
+    val_end = int(T * (train_ratio + val_ratio))
+
+    train_end = max(0, min(train_end, T))
+    val_end = max(train_end, min(val_end, T))
+
+    val_start = max(0, train_end - seq_len)
+    test_start = max(0, val_end - seq_len)
+
+    train_data = data[:train_end]
+    val_data = data[val_start:val_end]
+    test_data = data[test_start:]
+
+    return train_data, val_data, test_data
+
+
+def _fit_transform_splits(train_data, val_data, test_data, type_flag: str, scaler=None):
+    if type_flag == "1":
+        if scaler is None:
+            scaler = StandardScaler()
+            scaler.fit(train_data)
+        train_data = scaler.transform(train_data)
+        val_data = scaler.transform(val_data)
+        test_data = scaler.transform(test_data)
+        return train_data, val_data, test_data, scaler
+    else:
+        return train_data, val_data, test_data, None
+
+
+def _to_float32(x: np.ndarray) -> np.ndarray:
+    return np.asarray(x, dtype=np.float32)
+
+
+def _clean_numeric_csv(df: pd.DataFrame) -> np.ndarray:
+    """
+    Keep only numeric columns, and drop common junk index columns.
+    """
+    drop_cols = [c for c in df.columns if str(c).lower().startswith("unnamed")]
+    if drop_cols:
+        df = df.drop(columns=drop_cols, errors="ignore")
+
+    num_df = df.select_dtypes(include=[np.number])
+
+    if num_df.shape[1] == 0:
+        raise ValueError("No numeric columns found in CSV after cleaning. Check your file format.")
+
+    num_df = num_df.dropna(axis=0, how="any")
+
+    return num_df.values.astype(np.float32)
+
+
+
+class _BaseTimeSeriesDataset(Dataset):
+
+    def __init__(self, flag, seq_len, pre_len):
+        assert flag in ["train", "val", "test"]
+        self.flag = flag
+        self.seq_len = int(seq_len)
+        self.pre_len = int(pre_len)
+        self.scaler = None
+        self.split = None  
+
+    def __getitem__(self, index):
+        s_begin = index
+        s_end = s_begin + self.seq_len
+        r_end = s_end + self.pre_len
+
+        x = self.split[s_begin:s_end]
+        y = self.split[s_end:r_end]
+        return x, y
+
+    def __len__(self):
+        if self.split is None:
+            return 0
+        return max(0, len(self.split) - self.seq_len - self.pre_len)
+
+
+class Dataset_Dhfm(_BaseTimeSeriesDataset):
+    def __init__(self, root_path, flag, seq_len, pre_len, type, train_ratio, val_ratio, scaler=None):
+        super().__init__(flag, seq_len, pre_len)
+        self.path = root_path
+
+        load_data = np.load(root_path) 
+        data = np.array(load_data).transpose() 
+        data = _to_float32(data)
+
+        train_data, val_data, test_data = _split_with_overlap(data, train_ratio, val_ratio, self.seq_len)
+        train_data, val_data, test_data, self.scaler = _fit_transform_splits(train_data, val_data, test_data, type, scaler)
+
+        if self.flag == "train":
+            self.split = train_data
+        elif self.flag == "val":
+            self.split = val_data
+        else:
+            self.split = test_data
+
+
+class Dataset_ECG(_BaseTimeSeriesDataset):
+    def __init__(self, root_path, flag, seq_len, pre_len, type, train_ratio, val_ratio, scaler=None):
+        super().__init__(flag, seq_len, pre_len)
+        self.path = root_path
+
+        df = pd.read_csv(root_path)
+        data = _clean_numeric_csv(df)  
+
+        train_data, val_data, test_data = _split_with_overlap(data, train_ratio, val_ratio, self.seq_len)
+        train_data, val_data, test_data, self.scaler = _fit_transform_splits(train_data, val_data, test_data, type, scaler)
+
+        if self.flag == "train":
+            self.split = train_data
+        elif self.flag == "val":
+            self.split = val_data
+        else:
+            self.split = test_data
+
+class Dataset_Solar(_BaseTimeSeriesDataset):
+    def __init__(self, root_path, flag, seq_len, pre_len, type, train_ratio, val_ratio, scaler=None):
+        super().__init__(flag, seq_len, pre_len)
+        self.path = root_path
+
+        files = os.listdir(root_path)
+        solar_data = []
+        time_data = None
+
+        for file in files:
+            full = os.path.join(root_path, file)
+            if os.path.isdir(full):
+                continue
+            if file.startswith("DA_"):
+                arr = pd.read_csv(full).values
+                raw_time = arr[:, 0:1]
+                if time_data is None:
+                    time_data = raw_time
+                raw_data = arr[:, 1:arr.shape[1]]
+                raw_data = raw_data.transpose()
+                solar_data.append(raw_data)
+
+        if len(solar_data) == 0 or time_data is None:
+            raise ValueError(f"No solar files found in {root_path} with prefix 'DA_'.")
+
+        solar_data = np.array(solar_data).squeeze(1).transpose()   # (T, N)
+        time_data = np.array(time_data)                             # (T, 1)
+        out = np.concatenate((time_data, solar_data), axis=1)       # (T, 1+N)
+
+        filtered = []
+        for item in out:
+            dt = datetime.datetime.strptime(item[0], "%m/%d/%y %H:%M")
+            if 8 <= dt.hour <= 17:
+                filtered.append(item[1:out.shape[1]-1])
+
+        data = _to_float32(np.array(filtered))
+
+        train_data, val_data, test_data = _split_with_overlap(data, train_ratio, val_ratio, self.seq_len)
+        train_data, val_data, test_data, self.scaler = _fit_transform_splits(train_data, val_data, test_data, type, scaler)
+
+        if self.flag == "train":
+            self.split = train_data
+        elif self.flag == "val":
+            self.split = val_data
+        else:
+            self.split = test_data
+
+
+class Dataset_Wiki(_BaseTimeSeriesDataset):
+    def __init__(self, root_path, flag, seq_len, pre_len, type, train_ratio, val_ratio, scaler=None):
+        super().__init__(flag, seq_len, pre_len)
+        self.path = root_path
+
+        df = pd.read_csv(root_path)
+
+        if df.shape[1] < 2:
+            raise ValueError("Wiki CSV must have at least 2 columns (time + features).")
+
+        df_feat = df.iloc[:, 1:]
+        data = _clean_numeric_csv(df_feat)  
+
+        train_data, val_data, test_data = _split_with_overlap(data, train_ratio, val_ratio, self.seq_len)
+        train_data, val_data, test_data, self.scaler = _fit_transform_splits(train_data, val_data, test_data, type, scaler)
+
+        if self.flag == "train":
+            self.split = train_data
+        elif self.flag == "val":
+            self.split = val_data
+        else:
+            self.split = test_data
+
+
+
+class Dataset_PEMS_BAY(_BaseTimeSeriesDataset):
+    def __init__(self, root_path, flag, seq_len, pre_len, type, train_ratio, val_ratio, scaler=None, fillna="ffill"):
+        super().__init__(flag, seq_len, pre_len)
+        self.path = root_path
+
+        obj = pd.read_hdf(root_path)
+
+        if isinstance(obj, pd.Series):
+            df = obj.to_frame()
+        elif isinstance(obj, pd.DataFrame):
+            df = obj
+        else:
+            df = pd.DataFrame(obj)
+
+        if fillna == "ffill":
+            df = df.ffill()
+            df = df.fillna(0.0)  
+        elif fillna == "zero":
+            df = df.fillna(0.0)
+        elif fillna == "drop":
+            df = df.dropna(axis=0, how="any")
+        elif fillna is None:
+            pass
+        else:
+            raise ValueError("fillna must be one of: 'ffill', 'zero', 'drop', or None")
+
+        data = df.values.astype(np.float32)
+
+        train_data, val_data, test_data = _split_with_overlap(data, train_ratio, val_ratio, self.seq_len)
+        train_data, val_data, test_data, self.scaler = _fit_transform_splits(train_data, val_data, test_data, type, scaler)
+
+        if self.flag == "train":
+            self.split = train_data
+        elif self.flag == "val":
+            self.split = val_data
+        else:
+            self.split = test_data

model/FourierGNN.py

@@ -0,0 +1,168 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class FGN(nn.Module):
+    def __init__(self, args,  pre_length, embed_size,
+                 feature_size, seq_length, hidden_size, hard_thresholding_fraction=1, hidden_size_factor=1, sparsity_threshold=0.01):
+        super().__init__()
+        self.embed_size = embed_size
+        self.hidden_size = hidden_size
+        self.number_frequency = 1
+        self.pre_length = pre_length
+        self.feature_size = feature_size
+        self.seq_length = seq_length
+        self.frequency_size = self.embed_size // self.number_frequency
+        self.hidden_size_factor = hidden_size_factor
+        self.sparsity_threshold = sparsity_threshold
+        self.hard_thresholding_fraction = hard_thresholding_fraction
+        self.scale = 0.02
+        self.embeddings = nn.Parameter(torch.randn(1, self.embed_size))
+        self.args = args
+
+        self.w1 = nn.Parameter(
+            self.scale * torch.randn(2, self.frequency_size, self.frequency_size * self.hidden_size_factor))
+        self.b1 = nn.Parameter(self.scale * torch.randn(2, self.frequency_size * self.hidden_size_factor))
+        self.w2 = nn.Parameter(
+            self.scale * torch.randn(2, self.frequency_size * self.hidden_size_factor, self.frequency_size))
+        self.b2 = nn.Parameter(self.scale * torch.randn(2, self.frequency_size))
+        self.w3 = nn.Parameter(
+            self.scale * torch.randn(2, self.frequency_size,
+                                     self.frequency_size * self.hidden_size_factor))
+        self.b3 = nn.Parameter(
+            self.scale * torch.randn(2, self.frequency_size * self.hidden_size_factor))
+        self.embeddings_10 = nn.Parameter(torch.randn(self.seq_length, 8))
+        self.fc = nn.Sequential(
+            nn.Linear(self.embed_size * 8, 64),
+            nn.LeakyReLU(),
+            nn.Linear(64, self.hidden_size),
+            nn.LeakyReLU(),
+            nn.Linear(self.hidden_size, self.pre_length)
+        )
+        self.to('cuda:0')
+
+    def tokenEmb(self, x):
+        x = x.unsqueeze(2)
+        y = self.embeddings
+        return x * y
+
+    # FourierGNN
+    def fourierGC(self, x, B, N, L):
+        o1_real = torch.zeros([B, (N*L)//2 + 1, self.frequency_size * self.hidden_size_factor],
+                              device=x.device)
+        o1_imag = torch.zeros([B, (N*L)//2 + 1, self.frequency_size * self.hidden_size_factor],
+                              device=x.device)
+        o2_real = torch.zeros(x.shape, device=x.device)
+        o2_imag = torch.zeros(x.shape, device=x.device)
+
+        o3_real = torch.zeros(x.shape, device=x.device)
+        o3_imag = torch.zeros(x.shape, device=x.device)
+
+        o1_real = F.relu(
+            torch.einsum('bli,ii->bli', x.real, self.w1[0]) - \
+            torch.einsum('bli,ii->bli', x.imag, self.w1[1]) + \
+            self.b1[0]
+        )
+
+        o1_imag = F.relu(
+            torch.einsum('bli,ii->bli', x.imag, self.w1[0]) + \
+            torch.einsum('bli,ii->bli', x.real, self.w1[1]) + \
+            self.b1[1]
+        )
+
+        # 1 layer
+        y = torch.stack([o1_real, o1_imag], dim=-1)
+        y = F.softshrink(y, lambd=self.sparsity_threshold)
+
+        o2_real = F.relu(
+            torch.einsum('bli,ii->bli', o1_real, self.w2[0]) - \
+            torch.einsum('bli,ii->bli', o1_imag, self.w2[1]) + \
+            self.b2[0]
+        )
+
+        o2_imag = F.relu(
+            torch.einsum('bli,ii->bli', o1_imag, self.w2[0]) + \
+            torch.einsum('bli,ii->bli', o1_real, self.w2[1]) + \
+            self.b2[1]
+        )
+
+        # 2 layer
+        x = torch.stack([o2_real, o2_imag], dim=-1)
+        x = F.softshrink(x, lambd=self.sparsity_threshold)
+        x = x + y
+
+        o3_real = F.relu(
+                torch.einsum('bli,ii->bli', o2_real, self.w3[0]) - \
+                torch.einsum('bli,ii->bli', o2_imag, self.w3[1]) + \
+                self.b3[0]
+        )
+
+        o3_imag = F.relu(
+                torch.einsum('bli,ii->bli', o2_imag, self.w3[0]) + \
+                torch.einsum('bli,ii->bli', o2_real, self.w3[1]) + \
+                self.b3[1]
+        )
+
+        # 3 layer
+        z = torch.stack([o3_real, o3_imag], dim=-1)
+        z = F.softshrink(z, lambd=self.sparsity_threshold)
+        z = z + x
+        z = torch.view_as_complex(z)
+        return z
+
+    def forward(self, x):
+
+        if self.args.normalize:
+
+            mean = x.mean(dim=1, keepdim=True).detach()
+            x = x - mean
+
+            std = torch.sqrt(torch.var(x, dim=1, keepdim=True, unbiased=False) + 1e-5).detach()
+            x = x / std
+
+
+        x = x.permute(0, 2, 1).contiguous()
+        B, N, L = x.shape
+        # B*N*L ==> B*NL
+        x = x.reshape(B, -1)
+        # embedding B*NL ==> B*NL*D
+        x = self.tokenEmb(x)
+   
+
+        # FFT B*NL*D ==> B*NT/2*D
+        x = torch.fft.rfft(x, dim=1, norm='ortho')
+
+        x = x.reshape(B, (N*L)//2+1, self.frequency_size)
+
+        bias = x
+
+        # FourierGNN
+        x = self.fourierGC(x, B, N, L)
+
+        x = x + bias
+
+        x = x.reshape(B, (N*L)//2+1, self.embed_size)
+
+        # ifft
+        x = torch.fft.irfft(x, n=N*L, dim=1, norm="ortho")
+
+        x = x.reshape(B, N, L, self.embed_size)
+        x = x.permute(0, 1, 3, 2)  # B, N, D, L
+
+        # projection
+        x = torch.matmul(x, self.embeddings_10)
+        x = x.reshape(B, N, -1)
+        x = self.fc(x)
+        x = x.permute(0, 2, 1)
+
+        if self.args.normalize:
+            x = x * std
+            x = x + mean
+
+
+        aux = {
+            'gate_l0': torch.tensor(0.0, device=x.device) # placeholder
+        }
+
+        return x, aux
+

model/SpikF.py

@@ -0,0 +1,151 @@
+import torch
+from torch import nn
+from spikingjelly.clock_driven.neuron import MultiStepLIFNode
+
+class SPE(nn.Module):
+    def __init__(self, input_len, patch_num, patch_dim, T, tau, D):
+        super().__init__()
+        self.patch_projector = nn.Linear(input_len // patch_num, patch_dim)
+        self.bn = nn.BatchNorm2d(patch_dim)
+        self.encoder_lif = MultiStepLIFNode(tau=tau, detach_reset=False, backend='torch')
+
+        self.D = D
+        self.T = T
+        self.patch_dim = patch_dim
+        self.patch_num = patch_num
+
+    def forward(self, x):
+        B, L, D = x.shape
+
+        x = x.view(B, self.patch_num, L // self.patch_num, D).contiguous()
+        x = x.transpose(-1, -2).contiguous()
+        x = self.patch_projector(x)
+        x = x.repeat(self.T, 1, 1, 1, 1)
+        x = x.permute(0, 1, 4, 2, 3).contiguous()
+        x = x.flatten(0, 1)
+        x = self.bn(x)
+        x = x.view(self.T, B, self.patch_dim, self.patch_num, D)
+        x = self.encoder_lif(x)
+
+        return x
+
+class SFS(nn.Module):
+    def __init__(self, patch_num, D, patch_dim, tau, alpha):
+        super().__init__()
+        self.time2freq = nn.Linear(patch_num, patch_num // 2 + 1)
+
+        self.intra_conv = nn.Conv2d(in_channels=patch_dim, out_channels=patch_dim, kernel_size=[5, 1], stride=[1, 1], padding=[2, 0])
+        self.inter_conv = nn.Conv2d(in_channels=patch_dim, out_channels=patch_dim, kernel_size=[3, 1], stride=[1, 1], padding=[1, 0])
+
+        self.generator_lif = MultiStepLIFNode(tau=tau, detach_reset=True, backend='torch', v_threshold=0.1)
+        self.mp_lif = MultiStepLIFNode(tau=tau, detach_reset=True, backend='torch')
+        self.sfs_lif = MultiStepLIFNode(tau=tau, detach_reset=True, backend='torch')
+        self.intra_lif = MultiStepLIFNode(tau=tau, detach_reset=True, backend='torch')
+        self.inter_lif = MultiStepLIFNode(tau=tau, detach_reset=True, backend='torch')
+
+        self.bn1 = nn.BatchNorm2d(patch_dim)
+        self.bn2 = nn.BatchNorm2d(patch_dim)
+        self.bn3 = nn.BatchNorm2d(patch_dim)
+        self.bn4 = nn.BatchNorm2d(patch_dim)
+
+    def forward(self, x):
+        res_x = x
+        T, B, pd, pn, D = x.shape
+
+        x = x.transpose(-1, -2).contiguous()
+        freq_spec = torch.fft.rfft(x)
+
+        selector = self.time2freq(x)
+        selector = selector.flatten(0, 1)
+        selector = self.bn1(selector)
+        selector = selector.view(T, B, pd, D, -1)
+        selector = self.generator_lif(selector)
+        selector = selector.sum(dim=0, keepdim=True)
+        selector = self.mp_lif(selector)
+        selector = selector.repeat(T, 1, 1, 1, 1).float()
+        selector_imag = torch.zeros(selector.size()).to(x.device)
+        selector = torch.complex(selector, selector_imag).to(x.device)
+
+        remain_freq = selector * freq_spec
+
+        current = torch.fft.irfft(remain_freq)
+        current = current.transpose(-1, -2).contiguous()
+        current = current.flatten(0, 1)
+        current = self.bn2(current)
+        current = current.view(T, B, pd, pn, D)
+
+        spike = self.sfs_lif(current)
+        x = spike + res_x
+        res_x = x
+
+        x = x.flatten(0, 1)
+        x = self.intra_conv(x)
+        x = self.bn3(x)
+        x = x.view(T, B, pd, pn, D)
+        x = self.intra_lif(x) + res_x
+        res_x = x
+
+        x = x.transpose(0, 3).contiguous()
+        x = x.flatten(0, 1)
+        x = self.inter_conv(x)
+        x = self.bn4(x)
+        x = x.view(pn, B, pd, T, D)
+        x = x.transpose(0, 3)
+        x = self.inter_lif(x)
+        x = x + res_x
+        
+        return x
+
+class SpikF(nn.Module):
+    def __init__(self, args, input_len, patch_num, patch_dim, T, blocks, D, pred_len, tau, alpha, hidden_dim):
+        super().__init__()
+        self.SPE = SPE(input_len, patch_num, patch_dim, T, tau, D)
+        self.args = args
+
+        self.SFSs = nn.ModuleList()
+        for i in range(blocks):
+            self.SFSs.append(SFS(patch_num, D, patch_dim, tau, alpha))
+
+        self.dense1 = nn.Linear(patch_num * patch_dim, hidden_dim)
+        self.dense2 = nn.Linear(hidden_dim, pred_len)
+
+        self.bn = nn.BatchNorm1d(D)
+
+        self.activ = nn.GELU()
+
+    def forward(self, x):
+
+        if self.args.normalize:
+
+            mean = x.mean(dim=1, keepdim=True).detach()
+            x = x - mean
+
+            std = torch.sqrt(torch.var(x, dim=1, keepdim=True, unbiased=False) + 1e-5).detach()
+            x = x / std
+
+        x = self.SPE(x)
+        T, B, pd, pn, D = x.shape
+
+        for i in range(len(self.SFSs)):
+            x = self.SFSs[i](x)
+            
+        x = x.permute(0, 1, 4, 2, 3).contiguous()
+        x = x.flatten(-2, -1)
+        x = self.dense1(x)
+        x = x.flatten(0, 1)
+        x = self.bn(x)
+        x = self.activ(x)
+        x = self.dense2(x)
+        x = x.transpose(-1, -2).contiguous()
+        x = x.view(T, B, -1, D)
+
+        if self.args.normalize:
+            x = x * std
+            x = x + mean.repeat(T, 1, 1, 1)
+
+        aux = {
+            'gate_l0': torch.tensor(0.0, device=x.device)  # placeholder
+        }
+        
+
+        return x, aux

model/SpikF_GO.py

@@ -0,0 +1,445 @@
+from __future__ import annotations
+
+from typing import Dict, Tuple
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import weight_norm
+
+from spikingjelly.clock_driven.neuron import MultiStepLIFNode
+from spikingjelly.activation_based import surrogate
+
+
+
+
+class Affine(nn.Module):
+    def __init__(self, D: int):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.ones(D))
+        self.beta  = nn.Parameter(torch.zeros(D))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.gamma + self.beta
+
+
+
+class RMSNorm(nn.Module):
+    """
+    tok: [B, M, E]
+    Normalize over M per sample, per channel plus affine.
+    """
+    def __init__(self, E: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.affine = Affine(E)
+
+    def forward(self, tok: torch.Tensor) -> torch.Tensor:
+        rms = torch.rsqrt(tok.pow(2).mean(dim=1, keepdim=True) + self.eps)  # [B,1,E]
+        y = tok * rms
+        y = self.affine(y)
+        return y
+
+
+
+class SFFT(nn.Module):
+    """
+    S-FFT: implementing FFT on GPU; for theoretical information (spiking FFT),
+    refer to the our paper and paper SpikF.
+    """
+    def __init__(self, M: int):
+        super().__init__()
+        self.M = M
+        self.F = M // 2 + 1
+
+    def rfft(self, s_t: torch.Tensor) -> torch.Tensor:
+        T, B, M, E = s_t.shape
+        x = s_t.permute(0, 1, 3, 2).contiguous().view(T * B * E, M)  # [T*B*E, M]
+        Z = torch.fft.rfft(x, n=self.M, dim=-1, norm="ortho")        # [T*B*E, F] complex
+        Z = Z.view(T, B, E, self.F).permute(0, 1, 3, 2).contiguous() # [T,B,F,E]
+        return Z
+
+    def irfft(self, Z_t: torch.Tensor) -> torch.Tensor:
+        T, B, Freq, E = Z_t.shape
+        x = Z_t.permute(0, 1, 3, 2).contiguous().view(T * B * E, Freq)  # [T*B*E, F]
+        y = torch.fft.irfft(x, n=self.M, dim=-1, norm="ortho")          # [T*B*E, M]
+        y = y.view(T, B, E, self.M).permute(0, 1, 3, 2).contiguous()    # [T,B,M,E]
+        return y
+
+
+
+class HardConcreteGate(nn.Module):
+    """
+    Gate over frequency bins.
+    Z: [T,B,F,E] 
+    mask m: [1,1,F,1] in [0,1]
+    """
+    def __init__(self, F_bins: int, init_logit: float = 2.0, eps: float = 1e-6):
+        super().__init__()
+        self.log_alpha = nn.Parameter(torch.full((F_bins,), float(init_logit)))
+        self.eps = eps
+
+    def _sample_u(self, shape, device):
+        return torch.empty(shape, device=device).uniform_(self.eps, 1.0 - self.eps)
+
+    def _hard_concrete(self, training: bool, device, tau: float):
+        if training:
+            u = self._sample_u(self.log_alpha.shape, device)
+            s = torch.sigmoid((torch.log(u) - torch.log(1 - u) + self.log_alpha) / tau)
+        else:
+            s = torch.sigmoid(self.log_alpha)
+        s_bar = s * 1.2 - 0.1
+        return s_bar.clamp(0.0, 1.0)
+
+    def forward(self, Z: torch.Tensor, tau: float) -> Tuple[torch.Tensor, torch.Tensor]:
+        m = self._hard_concrete(self.training, Z.device, tau=tau)  # [F]
+        m = m.view(1, 1, -1, 1).to(Z.real.dtype)                   # [1,1,F,1]
+        return Z * m, m
+
+    def l0(self) -> torch.Tensor:
+        return torch.sigmoid(self.log_alpha).mean()
+
+
+
+
+class ComplexAffine(nn.Module):
+    def __init__(self, E: int):
+        super().__init__()
+        self.gamma_r = nn.Parameter(torch.ones(E))
+        self.beta_r  = nn.Parameter(torch.zeros(E))
+        self.gamma_i = nn.Parameter(torch.ones(E))
+        self.beta_i  = nn.Parameter(torch.zeros(E))
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        zr = z.real * self.gamma_r + self.beta_r
+        zi = z.imag * self.gamma_i + self.beta_i
+        return torch.complex(zr, zi)
+
+
+
+class ComplexLinear(nn.Module):
+    def __init__(self, E_in: int, E_out: int, init_scale: float = 0.02):
+        super().__init__()
+        self.Wr = nn.Parameter(init_scale * torch.randn(E_in, E_out))
+        self.Wi = nn.Parameter(init_scale * torch.randn(E_in, E_out))
+        self.br = nn.Parameter(torch.zeros(E_out))
+        self.bi = nn.Parameter(torch.zeros(E_out))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        xr, xi = x.real, x.imag
+        yr = xr @ self.Wr - xi @ self.Wi + self.br
+        yi = xi @ self.Wr + xr @ self.Wi + self.bi
+        return torch.complex(yr, yi)
+
+
+class ComplexLIFGate(nn.Module):
+    def __init__(self, tau: float, v_th: float):
+        super().__init__()
+        self.lif_r = MultiStepLIFNode(
+            tau=tau, v_threshold=v_th, detach_reset=True,
+            surrogate_function=surrogate.ATan(alpha=4.0), backend="torch"
+        )
+        self.lif_i = MultiStepLIFNode(
+            tau=tau, v_threshold=v_th, detach_reset=True,
+            surrogate_function=surrogate.ATan(alpha=4.0), backend="torch"
+        )
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        s_r = self.lif_r(z.real)  # [T,B,F,D] in [0,1]
+        s_i = self.lif_i(z.imag)
+        g = ((s_r > 0) | (s_i > 0)).to(z.real.dtype)
+        return g
+
+
+
+class SFGO(nn.Module):
+    def __init__(
+        self,
+        args,
+        E: int,
+        hidden_size_factor: int,
+        tau: float = 2.0,
+        v_th: float = 1.0,
+        apply_gate_to_complex: bool = True,
+    ):
+        super().__init__()
+        H = int(E * hidden_size_factor)
+
+        self.args = args
+
+        self.lin1 = ComplexLinear(E, H)
+        self.lin2 = ComplexLinear(H, E)
+        self.lin3 = ComplexLinear(E, E)
+
+        self.g1 = ComplexLIFGate(tau=tau, v_th=v_th)   
+        self.g2 = ComplexLIFGate(tau=tau, v_th=v_th)   
+        self.g3 = ComplexLIFGate(tau=tau, v_th=v_th)   
+
+        self.apply_gate_to_complex = apply_gate_to_complex
+
+        self.r2 = nn.Parameter(torch.tensor(0.1))
+        self.r3 = nn.Parameter(torch.tensor(0.1))
+
+        if self.args.affine:
+
+            self.a1 = ComplexAffine(E)
+            self.a2 = ComplexAffine(H)
+            self.a3 = ComplexAffine(E)
+
+            self.ga1 = ComplexLIFGate(tau=tau, v_th=v_th)  
+            self.ga2 = ComplexLIFGate(tau=tau, v_th=v_th)  
+            self.ga3 = ComplexLIFGate(tau=tau, v_th=v_th)  
+
+
+    def _apply_gate(self, z: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
+        if not self.apply_gate_to_complex:
+            return z
+        return z * g.to(z.real.dtype)
+
+    def forward(self, Z: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        stats: Dict[str, torch.Tensor] = {}
+
+
+        if self.args.affine:        
+            A1 = self.a1(Z)
+            GA1 = self.ga1(A1)
+            A1 = self._apply_gate(A1, GA1)
+        else:
+            A1 = Z
+
+        Y = self.lin1(A1)
+        G1 = self.g1(Y)
+        Y = self._apply_gate(Y, G1)
+
+        if self.args.affine:
+            A2 = self.a2(Y)
+            GA2 = self.ga2(A2)
+            A2 = self._apply_gate(A2, GA2)
+        else:
+            A2 = Y
+
+        X = self.lin2(A2)
+        G2 = self.g2(X)
+        X = self._apply_gate(X, G2)
+
+        Z2 = Z + self.r2 * X
+
+        if self.args.affine:
+            A3 = self.a3(Z2)
+            GA3 = self.ga3(A3)
+            A3 = self._apply_gate(A3, GA3)
+        else:
+            A3 = Z2
+
+
+        W = self.lin3(A3)
+        G3 = self.g3(W)
+        W = self._apply_gate(W, G3)
+
+        out = Z2 + self.r3 * W
+
+        with torch.no_grad():
+            mag2 = out.real * out.real + out.imag * out.imag
+            stats["freq_active_frac"] = (mag2 > 0).float().mean()
+
+            stats["rezero_r2"] = self.r2.detach()
+            stats["rezero_r3"] = self.r3.detach()
+
+            stats["gate_lin_frac_1"] = G1.mean().detach()
+            stats["gate_lin_frac_2"] = G2.mean().detach()
+            stats["gate_lin_frac_3"] = G3.mean().detach()
+
+        return out, stats
+
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        E: int,
+        L: int,
+        pred_len: int,
+        T: int,
+        tau: float,
+        v_th: float,
+        proj_dim: int = 4,
+        reduced_dim: int = 64,
+    ):
+        super().__init__()
+        self.E, self.L, self.P, self.T = E, L, pred_len, T
+        self.proj_dim = int(proj_dim)
+
+        self.time_proj = nn.Linear(L, self.proj_dim, bias=False)
+        D_in = E * self.proj_dim
+        self.reduced_dim = int(reduced_dim)
+
+        self.lif = MultiStepLIFNode(
+            tau=tau,
+            v_threshold=v_th,
+            detach_reset=True,
+            surrogate_function=surrogate.ATan(alpha=4.0),
+            backend="torch",
+        )
+
+        self.fc_reduce = weight_norm(nn.Linear(D_in, int(reduced_dim), bias=True))
+        self.fc_out    = weight_norm(nn.Linear(int(reduced_dim), pred_len, bias=True))
+
+        nn.init.xavier_uniform_(self.time_proj.weight, gain=0.5)
+        nn.init.xavier_uniform_(self.fc_reduce.weight, gain=0.6)
+        nn.init.xavier_uniform_(self.fc_out.weight, gain=0.2)
+        nn.init.zeros_(self.fc_reduce.bias)
+        nn.init.zeros_(self.fc_out.bias)
+
+    def forward(self, y_t: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        T, B, N, E, L = y_t.shape
+
+        y_p = self.time_proj(y_t)                        # [T,B,N,E,p]
+        x   = y_p.reshape(T, B * N, E * self.proj_dim)   # [T,B*N,D]
+        s   = self.lif(x)                                # [T,B*N,D] spikes
+        h_t = self.fc_reduce(s.reshape(T * B * N, -1)).view(T, B * N, self.reduced_dim)
+
+        h = h_t.mean(dim=0)                           # [B*N,reduced_dim]
+        h = F.gelu(h)
+        out = self.fc_out(h)                             # [B*N,O]
+
+        preds = out.view(B, N, self.P).permute(0, 2, 1).contiguous()
+        stats = {"dec_spike_rate": s.mean().detach()}
+        return preds, stats
+
+
+
+class SpikF_GO(nn.Module):
+    def __init__(
+        self,
+        args,
+        pre_length: int,
+        embed_size: int,
+        feature_size: int,
+        seq_length: int,
+        hidden_size: int,
+        hard_thresholding_fraction=1,
+        hidden_size_factor: int = 1,
+        sparsity_threshold: float = 0.01,
+    ):
+        super().__init__()
+        self.args = args
+
+        self.N = feature_size
+        self.L = seq_length
+        self.E = embed_size
+        self.T = args.T
+        self.M = self.N * self.L
+
+
+        self.embeddings = nn.Parameter(torch.randn(1, self.E) * 0.02)
+        self.node_aff = Affine(self.E)
+        self.node_rms = RMSNorm(E=self.E, eps=1e-6)
+
+        # step modulation
+        self.step_gamma = nn.Parameter(torch.ones(self.T))
+        self.step_beta  = nn.Parameter(torch.zeros(self.T))
+        self.register_buffer("step_scale", torch.linspace(0, 1, steps=self.T).view(self.T, 1, 1, 1))
+
+        # Encoder LIF
+        self.encoder_lif = MultiStepLIFNode(
+            tau=args.tau,
+            v_threshold=args.alpha,
+            detach_reset=True,
+            surrogate_function=surrogate.ATan(alpha=4.0),
+            backend="torch",
+        )
+
+        self.sfft = SFFT(self.M)
+        self.F_bins = self.sfft.F
+
+        # frequency gate
+        self.freq_gate = HardConcreteGate(self.F_bins, init_logit=2.0)
+        self.register_buffer("gate_tau", torch.tensor(0.10))
+
+        self.sfgo = SFGO(
+            self.args,
+            E=self.E,
+            hidden_size_factor=hidden_size_factor,
+            tau=args.tau,
+            v_th=args.alpha,
+            apply_gate_to_complex=True,
+        )
+
+        # decoder
+        proj_dim = self.args.proj_dim
+        reduced_dim = max(16, min(128, hidden_size // 4))
+        self.decoder = Decoder(
+            E=self.E,
+            L=self.L,
+            pred_len=pre_length,
+            T=self.T,
+            tau=args.tau,
+            v_th=args.alpha,
+            proj_dim=proj_dim,
+            reduced_dim=reduced_dim,
+        )
+
+    def node_embed(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B,L,N] -> [B,M,E]
+        B, L, N = x.shape
+        x_flat = x.permute(0, 2, 1).contiguous().reshape(B, self.M)  # [B,M]
+        tok = x_flat.unsqueeze(-1) * self.embeddings                 # [B,M,E]
+        tok = self.node_aff(tok)
+        return tok
+        
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        B, L, N = x.shape
+
+        # normalize
+        if self.args.normalize:
+            mean = x.mean(dim=1, keepdim=True).detach()
+            x0 = x - mean
+            std = torch.sqrt(torch.var(x0, dim=1, keepdim=True, unbiased=False) + 1e-5).detach()
+            x0 = x0 / std
+        else:
+            mean, std = None, None
+            x0 = x
+
+
+        tok = self.node_embed(x0)         # [B,M,E]
+        tok = self.node_rms(tok)            # RMSNorm
+
+
+        # step modulation
+        cur_t = tok.unsqueeze(0).repeat(self.T, 1, 1, 1)
+        cur_t = cur_t * self.step_gamma.view(self.T, 1, 1, 1) + self.step_beta.view(self.T, 1, 1, 1)
+        cur_t = cur_t * (1.0 + 0.02 * self.step_scale.to(cur_t.dtype))
+
+
+        # spikes
+        s_t = self.encoder_lif(cur_t)
+        enc_rate = s_t.mean()
+
+        # FFT
+        Z_t = self.sfft.rfft(s_t)
+
+        # prune
+        Z_t, m = self.freq_gate(Z_t, tau=float(self.gate_tau))
+
+        # S-FGO blocks
+        Z_t, fb_stats = self.sfgo(Z_t)
+
+        # iFFT
+        y_time_t = self.sfft.irfft(Z_t).to(tok.dtype)
+
+        y_t = y_time_t.view(self.T, B, N, self.L, self.E).permute(0, 1, 2, 4, 3).contiguous()
+
+        preds, dec_stats = self.decoder(y_t)
+
+        if self.args.normalize:
+            preds = preds * std + mean # denormalize
+
+        aux = {
+            "enc_rate": enc_rate.detach(),
+            "rho_hat": self.freq_gate.l0().detach(),
+            "freq_mask_mean": m.mean().detach(),
+            "freq_mask_active": (m > 0.5).float().mean().detach(),
+            **fb_stats,
+            **dec_stats,
+        }
+        return preds, aux

model/SpikF_GO_CPG.py

@@ -0,0 +1,514 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Tuple, Dict, Optional
+import torch
+import math
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import weight_norm
+
+from spikingjelly.clock_driven.neuron import MultiStepLIFNode
+from spikingjelly.activation_based import surrogate
+
+
+
+
+class Affine(nn.Module):
+    def __init__(self, D: int):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.ones(D))
+        self.beta  = nn.Parameter(torch.zeros(D))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.gamma + self.beta
+
+
+
+class RMSNorm(nn.Module):
+    """
+    tok: [B, M, E]
+    Normalize over M per sample, per channel plus affine.
+    """
+    def __init__(self, E: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.affine = Affine(E)
+
+    def forward(self, tok: torch.Tensor) -> torch.Tensor:
+        rms = torch.rsqrt(tok.pow(2).mean(dim=1, keepdim=True) + self.eps)  # [B,1,E]
+        y = tok * rms
+        y = self.affine(y)
+        return y
+
+
+
+
+class CPGSpikePE(nn.Module):
+    """
+    Spike-form positional encoding (CPG-PE).
+    Generates 2*N_pe binary channels with log-spaced rhythms over the flattened index t in [0, T*M).
+    Shapes:
+      returns pe: [T, B, M, 2*N_pe] with 0/1 spikes (no learnable params).
+    """
+    def __init__(self,
+                 num_pairs: int = 20,
+                 tau: float = 10000.0,
+                 eta: float = 1.0,
+                 vthres: float = 0.8,
+                 w_max: float = 10000.0):
+        super().__init__()
+        self.num_pairs = num_pairs
+        self.tau = tau
+        self.eta = eta
+        self.vthres = vthres
+        self.w_max = w_max
+
+    def forward(self, T: int, B: int, M: int, device) -> torch.Tensor:
+        t = torch.arange(T * M, device=device, dtype=torch.float32)  # [T*M]
+        i = torch.arange(self.num_pairs, device=device, dtype=torch.float32)
+        freq = torch.exp(-torch.log(torch.tensor(self.w_max, device=device)) * (i / max(1, self.num_pairs)))  # [N_pe]
+
+        arg = self.eta * (t[:, None] * freq[None, :] / self.tau)  # [T*M, N_pe]
+        cos_spk = (torch.cos(arg) - self.vthres > 0).float()
+        sin_spk = (torch.sin(arg) - self.vthres > 0).float()
+
+        pe = torch.cat([cos_spk, sin_spk], dim=1)                      # [T*M, 2*N_pe]
+        pe = pe.view(T, M, 2 * self.num_pairs).unsqueeze(1)            # [T, 1, M, 2*N_pe]
+        pe = pe.expand(-1, B, -1, -1).contiguous()                     # [T, B, M, 2*N_pe]
+        return pe
+
+
+
+
+class SFFT(nn.Module):
+    """
+    S-FFT: implementing FFT on GPU; for theoretical information (spiking FFT),
+    refer to the our paper and paper SpikF.
+    """
+    def __init__(self, M: int):
+        super().__init__()
+        self.M = M
+        self.F = M // 2 + 1
+
+    def rfft(self, s_t: torch.Tensor) -> torch.Tensor:
+        T, B, M, E = s_t.shape
+        x = s_t.permute(0, 1, 3, 2).contiguous().view(T * B * E, M)  # [T*B*E, M]
+        Z = torch.fft.rfft(x, n=self.M, dim=-1, norm="ortho")        # [T*B*E, F] complex
+        Z = Z.view(T, B, E, self.F).permute(0, 1, 3, 2).contiguous() # [T,B,F,E]
+        return Z
+
+    def irfft(self, Z_t: torch.Tensor) -> torch.Tensor:
+        T, B, Freq, E = Z_t.shape
+        x = Z_t.permute(0, 1, 3, 2).contiguous().view(T * B * E, Freq)  # [T*B*E, F]
+        y = torch.fft.irfft(x, n=self.M, dim=-1, norm="ortho")          # [T*B*E, M]
+        y = y.view(T, B, E, self.M).permute(0, 1, 3, 2).contiguous()    # [T,B,M,E]
+        return y
+
+
+
+class HardConcreteGate(nn.Module):
+    """
+    Gate over frequency bins.
+    Z: [T,B,F,E] 
+    mask m: [1,1,F,1] in [0,1]
+    """
+    def __init__(self, F_bins: int, init_logit: float = 2.0, eps: float = 1e-6):
+        super().__init__()
+        self.log_alpha = nn.Parameter(torch.full((F_bins,), float(init_logit)))
+        self.eps = eps
+
+    def _sample_u(self, shape, device):
+        return torch.empty(shape, device=device).uniform_(self.eps, 1.0 - self.eps)
+
+    def _hard_concrete(self, training: bool, device, tau: float):
+        if training:
+            u = self._sample_u(self.log_alpha.shape, device)
+            s = torch.sigmoid((torch.log(u) - torch.log(1 - u) + self.log_alpha) / tau)
+        else:
+            s = torch.sigmoid(self.log_alpha)
+        s_bar = s * 1.2 - 0.1
+        return s_bar.clamp(0.0, 1.0)
+
+    def forward(self, Z: torch.Tensor, tau: float) -> Tuple[torch.Tensor, torch.Tensor]:
+        m = self._hard_concrete(self.training, Z.device, tau=tau)  # [F]
+        m = m.view(1, 1, -1, 1).to(Z.real.dtype)                   # [1,1,F,1]
+        return Z * m, m
+
+    def l0(self) -> torch.Tensor:
+        return torch.sigmoid(self.log_alpha).mean()
+
+
+
+
+class ComplexAffine(nn.Module):
+    def __init__(self, E: int):
+        super().__init__()
+        self.gamma_r = nn.Parameter(torch.ones(E))
+        self.beta_r  = nn.Parameter(torch.zeros(E))
+        self.gamma_i = nn.Parameter(torch.ones(E))
+        self.beta_i  = nn.Parameter(torch.zeros(E))
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        zr = z.real * self.gamma_r + self.beta_r
+        zi = z.imag * self.gamma_i + self.beta_i
+        return torch.complex(zr, zi)
+
+
+
+class ComplexLinear(nn.Module):
+    def __init__(self, E_in: int, E_out: int, init_scale: float = 0.02):
+        super().__init__()
+        self.Wr = nn.Parameter(init_scale * torch.randn(E_in, E_out))
+        self.Wi = nn.Parameter(init_scale * torch.randn(E_in, E_out))
+        self.br = nn.Parameter(torch.zeros(E_out))
+        self.bi = nn.Parameter(torch.zeros(E_out))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        xr, xi = x.real, x.imag
+        yr = xr @ self.Wr - xi @ self.Wi + self.br
+        yi = xi @ self.Wr + xr @ self.Wi + self.bi
+        return torch.complex(yr, yi)
+
+
+class ComplexLIFGate(nn.Module):
+    def __init__(self, tau: float, v_th: float):
+        super().__init__()
+        self.lif_r = MultiStepLIFNode(
+            tau=tau, v_threshold=v_th, detach_reset=True,
+            surrogate_function=surrogate.ATan(alpha=4.0), backend="torch"
+        )
+        self.lif_i = MultiStepLIFNode(
+            tau=tau, v_threshold=v_th, detach_reset=True,
+            surrogate_function=surrogate.ATan(alpha=4.0), backend="torch"
+        )
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        s_r = self.lif_r(z.real)  # [T,B,F,D] in [0,1]
+        s_i = self.lif_i(z.imag)
+        g = ((s_r > 0) | (s_i > 0)).to(z.real.dtype)
+        return g
+
+
+
+class SFGO(nn.Module):
+    def __init__(
+        self,
+        args,
+        E: int,
+        hidden_size_factor: int,
+        tau: float = 2.0,
+        v_th: float = 1.0,
+        apply_gate_to_complex: bool = True,
+    ):
+        super().__init__()
+        H = int(E * hidden_size_factor)
+
+        self.args = args
+
+        self.lin1 = ComplexLinear(E, H)
+        self.lin2 = ComplexLinear(H, E)
+        self.lin3 = ComplexLinear(E, E)
+
+        self.g1 = ComplexLIFGate(tau=tau, v_th=v_th)   
+        self.g2 = ComplexLIFGate(tau=tau, v_th=v_th)   
+        self.g3 = ComplexLIFGate(tau=tau, v_th=v_th)   
+
+        self.apply_gate_to_complex = apply_gate_to_complex
+
+        self.r2 = nn.Parameter(torch.tensor(0.1))
+        self.r3 = nn.Parameter(torch.tensor(0.1))
+
+        if self.args.affine:
+
+            self.a1 = ComplexAffine(E)
+            self.a2 = ComplexAffine(H)
+            self.a3 = ComplexAffine(E)
+
+            self.ga1 = ComplexLIFGate(tau=tau, v_th=v_th)  
+            self.ga2 = ComplexLIFGate(tau=tau, v_th=v_th)  
+            self.ga3 = ComplexLIFGate(tau=tau, v_th=v_th)  
+
+
+    def _apply_gate(self, z: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
+        if not self.apply_gate_to_complex:
+            return z
+        return z * g.to(z.real.dtype)
+
+    def forward(self, Z: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        stats: Dict[str, torch.Tensor] = {}
+
+
+        if self.args.affine:        
+            A1 = self.a1(Z)
+            GA1 = self.ga1(A1)
+            A1 = self._apply_gate(A1, GA1)
+        else:
+            A1 = Z
+
+        Y = self.lin1(A1)
+        G1 = self.g1(Y)
+        Y = self._apply_gate(Y, G1)
+
+        if self.args.affine:
+            A2 = self.a2(Y)
+            GA2 = self.ga2(A2)
+            A2 = self._apply_gate(A2, GA2)
+        else:
+            A2 = Y
+
+        X = self.lin2(A2)
+        G2 = self.g2(X)
+        X = self._apply_gate(X, G2)
+
+        Z2 = Z + self.r2 * X
+
+        if self.args.affine:
+            A3 = self.a3(Z2)
+            GA3 = self.ga3(A3)
+            A3 = self._apply_gate(A3, GA3)
+        else:
+            A3 = Z2
+
+
+        W = self.lin3(A3)
+        G3 = self.g3(W)
+        W = self._apply_gate(W, G3)
+
+        out = Z2 + self.r3 * W
+
+        with torch.no_grad():
+            mag2 = out.real * out.real + out.imag * out.imag
+            stats["freq_active_frac"] = (mag2 > 0).float().mean()
+
+            stats["rezero_r2"] = self.r2.detach()
+            stats["rezero_r3"] = self.r3.detach()
+
+            stats["gate_lin_frac_1"] = G1.mean().detach()
+            stats["gate_lin_frac_2"] = G2.mean().detach()
+            stats["gate_lin_frac_3"] = G3.mean().detach()
+
+        return out, stats
+
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        E: int,
+        L: int,
+        pred_len: int,
+        T: int,
+        tau: float,
+        v_th: float,
+        proj_dim: int = 4,
+        reduced_dim: int = 64,
+    ):
+        super().__init__()
+        self.E, self.L, self.P, self.T = E, L, pred_len, T
+        self.proj_dim = int(proj_dim)
+
+        self.time_proj = nn.Linear(L, self.proj_dim, bias=False)
+        D_in = E * self.proj_dim
+        self.reduced_dim = int(reduced_dim)
+
+        self.lif = MultiStepLIFNode(
+            tau=tau,
+            v_threshold=v_th,
+            detach_reset=True,
+            surrogate_function=surrogate.ATan(alpha=4.0),
+            backend="torch",
+        )
+
+        self.fc_reduce = weight_norm(nn.Linear(D_in, int(reduced_dim), bias=True))
+        self.fc_out    = weight_norm(nn.Linear(int(reduced_dim), pred_len, bias=True))
+
+        nn.init.xavier_uniform_(self.time_proj.weight, gain=0.5)
+        nn.init.xavier_uniform_(self.fc_reduce.weight, gain=0.6)
+        nn.init.xavier_uniform_(self.fc_out.weight, gain=0.2)
+        nn.init.zeros_(self.fc_reduce.bias)
+        nn.init.zeros_(self.fc_out.bias)
+
+    def forward(self, y_t: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        T, B, N, E, L = y_t.shape
+
+        y_p = self.time_proj(y_t)                        # [T,B,N,E,p]
+        x   = y_p.reshape(T, B * N, E * self.proj_dim)   # [T,B*N,D]
+        s   = self.lif(x)                                # [T,B*N,D] spikes
+        h_t = self.fc_reduce(s.reshape(T * B * N, -1)).view(T, B * N, self.reduced_dim)
+
+        h = h_t.mean(dim=0)                           # [B*N,reduced_dim]
+        h = F.gelu(h)
+        out = self.fc_out(h)                             # [B*N,O]
+
+        preds = out.view(B, N, self.P).permute(0, 2, 1).contiguous()
+        stats = {"dec_spike_rate": s.mean().detach()}
+        return preds, stats
+
+
+
+class SpikF_GO_CPG(nn.Module):
+    def __init__(
+        self,
+        args,
+        pre_length: int,
+        embed_size: int,
+        feature_size: int,
+        seq_length: int,
+        hidden_size: int,
+        hard_thresholding_fraction=1,
+        hidden_size_factor: int = 1,
+        sparsity_threshold: float = 0.01,
+    ):
+        super().__init__()
+        self.args = args
+
+        self.N = feature_size
+        self.L = seq_length
+        self.E = embed_size
+        self.T = args.T
+        self.M = self.N * self.L
+
+        self.use_cpg_pe = True
+        self.num_pe_pairs = 20
+        self.pe_tau = 10000.0
+        self.pe_eta = 1.0
+        self.pe_vthres = 0.8
+        self.pe_wmax = 10000.0
+
+
+        if self.use_cpg_pe:
+            self.cpg_pe = CPGSpikePE(
+                num_pairs=self.num_pe_pairs,
+                tau=self.pe_tau, eta=self.pe_eta,
+                vthres=self.pe_vthres, w_max=self.pe_wmax
+            )
+            self.pe_linear = nn.Linear(self.E + 2 * self.num_pe_pairs, self.E, bias=False)
+            self.pe_bn = nn.BatchNorm1d(self.E)
+            self.pe_lif = MultiStepLIFNode(
+                tau=self.args.tau, v_threshold=self.args.alpha, detach_reset=True,
+                surrogate_function=surrogate.ATan(alpha=4.0), backend='torch'
+            )
+
+
+
+        self.embeddings = nn.Parameter(torch.randn(1, self.E) * 0.02)
+        self.node_aff = Affine(self.E)
+        self.node_rms = RMSNorm(E=self.E, eps=1e-6)
+
+        # step modulation
+        self.step_gamma = nn.Parameter(torch.ones(self.T))
+        self.step_beta  = nn.Parameter(torch.zeros(self.T))
+        self.register_buffer("step_scale", torch.linspace(0, 1, steps=self.T).view(self.T, 1, 1, 1))
+
+        # Encoder LIF
+        self.encoder_lif = MultiStepLIFNode(
+            tau=args.tau,
+            v_threshold=args.alpha,
+            detach_reset=True,
+            surrogate_function=surrogate.ATan(alpha=4.0),
+            backend="torch",
+        )
+
+        self.sfft = SFFT(self.M)
+        self.F_bins = self.sfft.F
+
+        # frequency gate
+        self.freq_gate = HardConcreteGate(self.F_bins, init_logit=2.0)
+        self.register_buffer("gate_tau", torch.tensor(0.10))
+
+        self.sfgo = SFGO(
+            self.args,
+            E=self.E,
+            hidden_size_factor=hidden_size_factor,
+            tau=args.tau,
+            v_th=args.alpha,
+            apply_gate_to_complex=True,
+        )
+
+        # decoder
+        proj_dim = self.args.proj_dim
+        reduced_dim = max(16, min(128, hidden_size // 4))
+        self.decoder = Decoder(
+            E=self.E,
+            L=self.L,
+            pred_len=pre_length,
+            T=self.T,
+            tau=args.tau,
+            v_th=args.alpha,
+            proj_dim=proj_dim,
+            reduced_dim=reduced_dim,
+        )
+
+    def node_embed(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B,L,N] -> [B,M,E]
+        B, L, N = x.shape
+        x_flat = x.permute(0, 2, 1).contiguous().reshape(B, self.M)  # [B,M]
+        tok = x_flat.unsqueeze(-1) * self.embeddings                 # [B,M,E]
+        tok = self.node_aff(tok)
+        return tok
+        
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        B, L, N = x.shape
+
+        # normalize
+        if self.args.normalize:
+            mean = x.mean(dim=1, keepdim=True).detach()
+            x0 = x - mean
+            std = torch.sqrt(torch.var(x0, dim=1, keepdim=True, unbiased=False) + 1e-5).detach()
+            x0 = x0 / std
+        else:
+            mean, std = None, None
+            x0 = x
+
+
+        tok = self.node_embed(x0)         # [B,M,E]
+        tok = self.node_rms(tok)            # RMSNorm
+
+
+        # step modulation
+        cur_t = tok.unsqueeze(0).repeat(self.T, 1, 1, 1)
+        cur_t = cur_t * self.step_gamma.view(self.T, 1, 1, 1) + self.step_beta.view(self.T, 1, 1, 1)
+        cur_t = cur_t * (1.0 + 0.02 * self.step_scale.to(cur_t.dtype))
+
+
+        # spikes
+        s_t = self.encoder_lif(cur_t)
+        if self.use_cpg_pe:
+            pe_spk = self.cpg_pe(T=self.T, B=B, M=self.M, device=x.device)      # [T,B,M,2*N_pe]
+            s_cat = torch.cat([s_t, pe_spk], dim=-1)                        # [T,B,M,E+2*N_pe]
+            h = self.pe_linear(s_cat)                                       # [T,B,M,E]
+            h = h.reshape(self.T * B * self.M, self.E)
+            h = self.pe_bn(h).view(self.T, B, self.M, self.E)
+            s_t = self.pe_lif(h)  
+
+        enc_rate = s_t.mean()
+
+        # FFT
+        Z_t = self.sfft.rfft(s_t)
+
+        # prune
+        Z_t, m = self.freq_gate(Z_t, tau=float(self.gate_tau))
+
+        # S-FGO blocks
+        Z_t, fb_stats = self.sfgo(Z_t)
+
+        # iFFT
+        y_time_t = self.sfft.irfft(Z_t).to(tok.dtype)
+
+        y_t = y_time_t.view(self.T, B, N, self.L, self.E).permute(0, 1, 2, 4, 3).contiguous()
+
+        preds, dec_stats = self.decoder(y_t)
+
+        if self.args.normalize:
+            preds = preds * std + mean # denormalize
+
+        aux = {
+            "enc_rate": enc_rate.detach(),
+            "rho_hat": self.freq_gate.l0().detach(),
+            "freq_mask_mean": m.mean().detach(),
+            "freq_mask_active": (m > 0.5).float().mean().detach(),
+            **fb_stats,
+            **dec_stats,
+        }
+        return preds, aux

model/SpikeGRU.py

@@ -0,0 +1,241 @@
+from typing import Optional
+from pathlib import Path
+
+from spikingjelly.activation_based import surrogate as sj_surrogate
+from snntorch import utils
+import snntorch as snn
+from snntorch import surrogate
+import torch
+from torch import nn
+
+
+
+class GRUCell(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        num_steps: int = 4,
+        grad_slope: float = 25.0,
+        beta: float = 0.99,
+        output_mems: bool = False,
+    ):
+        super().__init__()
+        self.spike_grad = surrogate.atan(alpha=2.0)
+        self.input_size = input_size
+        self.num_steps = num_steps
+        self.hidden_size = hidden_size
+        self.beta = beta
+        self.full_rec = output_mems
+        self.lif = snn.Leaky(
+            beta=self.beta,
+            spike_grad=self.spike_grad,
+            init_hidden=True,
+            output=output_mems,
+        )
+        self.linear_ih = nn.Linear(input_size, 3 * hidden_size)
+        self.linear_hh = nn.Linear(hidden_size, 3 * hidden_size)
+        self.surrogate_function1 = sj_surrogate.ATan()
+
+    def forward(self, inputs):
+        if inputs.size(-1) == self.input_size:
+            # assume static spikes:
+            h = torch.zeros(
+                size=[inputs.shape[0], self.hidden_size],
+                dtype=torch.float,
+                device=inputs.device,
+            )
+            y_ih = torch.split(self.linear_ih(inputs), self.hidden_size, dim=1)
+            y_hh = torch.split(self.linear_hh(h), self.hidden_size, dim=1)
+            r = self.surrogate_function1(y_ih[0] + y_hh[0])
+            z = self.surrogate_function1(y_ih[1] + y_hh[1])
+            n = self.surrogate_function1(y_ih[2] + r * y_hh[2])
+            h = (1.0 - z) * n + z * h
+            cur = h
+            static = True
+        elif inputs.size(-1) == self.num_steps and inputs.size(-2) == self.input_size:
+            inputs = inputs.transpose(-1, -2)  # BC, T, H
+            h = torch.zeros(
+                size=[inputs.shape[0], self.hidden_size, self.num_steps],
+                dtype=torch.float,
+                device=inputs.device,
+            )
+            y_ih = torch.split(
+                self.linear_ih(inputs).transpose(-1, -2), self.hidden_size, dim=1
+            )
+            y_hh = torch.split(
+                self.linear_hh(h.transpose(-1, -2)).transpose(-1, -2),
+                self.hidden_size,
+                dim=1,
+            )
+            r = self.surrogate_function1(y_ih[0] + y_hh[0])
+            z = self.surrogate_function1(y_ih[1] + y_hh[1])
+            n = self.surrogate_function1(y_ih[2] + r * y_hh[2])
+            h = (1.0 - z) * n + z * h
+            cur = h
+            static = False
+        else:
+            raise ValueError(
+                f"Input size mismatch!"
+                f"Got {inputs.size()} but expected (..., {self.input_size}, {self.num_steps}) or (..., {self.input_size})"
+            )
+
+        spk_rec = []
+        mem_rec = []
+        if self.full_rec:
+            for i_step in range(self.num_steps):
+                if static:
+                    spk, mem = self.lif(cur)
+                else:
+                    spk, mem = self.lif(cur[:, :, i_step])
+                spk_rec.append(spk)
+                mem_rec.append(mem)
+            spks = torch.stack(spk_rec, dim=-1)
+            mems = torch.stack(mem_rec, dim=-1)
+            return spks, mems
+        else:
+            for i_step in range(self.num_steps):
+                if static:
+                    spk = self.lif(cur)
+                else:
+                    spk = self.lif(cur[:, :, i_step])
+                spk_rec.append(spk)
+            spks = torch.stack(spk_rec, dim=-1)
+            return spks
+
+
+class DeltaEncoder(nn.Module):
+    def __init__(self, output_size: int):
+        super().__init__()
+        self.norm = nn.BatchNorm2d(1)
+        self.enc = nn.Linear(1, output_size)
+        self.lif = snn.Leaky(
+            beta=0.99, spike_grad=surrogate.atan(), init_hidden=True, output=False
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: batch, L, C
+        delta = torch.zeros_like(inputs)
+        delta[:, 1:] = inputs[:, 1:, :] - inputs[:, :-1, :]
+        delta = delta.unsqueeze(1).permute(0, 1, 3, 2)  # batch, 1, C, L
+        delta = self.norm(delta)
+        delta = delta.permute(0, 2, 3, 1)  # batch, C, L, 1
+        enc = self.enc(delta)  # batch, C, L, output_size
+        enc = enc.permute(0, 3, 1, 2)  # batch, output_size, C, L
+        spks = self.lif(enc)
+        return spks
+
+
+class ConvEncoder(nn.Module):
+    def __init__(self, output_size: int, kernel_size: int = 3):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv2d(
+                in_channels=1,
+                out_channels=output_size,
+                kernel_size=(1, kernel_size),
+                stride=1,
+                padding=(0, kernel_size // 2),
+            ),
+            nn.BatchNorm2d(output_size),
+        )
+        self.lif = snn.Leaky(
+            beta=0.99,
+            spike_grad=surrogate.atan(alpha=2.0),
+            init_hidden=True,
+            output=False,
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: batch, L, C
+        inputs = inputs.permute(0, 2, 1).unsqueeze(1)  # batch, 1, C, L
+        enc = self.encoder(inputs)  # batch, output_size, C, L
+        spks = self.lif(enc)
+        return spks
+
+
+
+class SpikeGRU(nn.Module):
+    def __init__(
+        self,
+        args,
+        hidden_size: int,
+        layers: int = 1,
+        num_steps: int = 50,
+        grad_slope: float = 25.0,
+        input_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        weight_file: Optional[Path] = None,
+        encoder_type: Optional[str] = "conv",
+    ):
+        super().__init__()
+        self.args = args
+        self.hidden_size   = args.hidden_size
+        self.num_steps   = args.T
+        self.input_size = args.feature_size
+        self.pre_length   = args.pre_length
+        self.layers       = args.blocks
+
+
+        if encoder_type == "conv":
+            self.encoder = ConvEncoder(self.hidden_size)
+        elif encoder_type == "delta":
+            self.encoder = DeltaEncoder(self.hidden_size)
+        else:
+            raise ValueError(f"Unknown encoder type {encoder_type}")
+
+        self.net = nn.Sequential(
+            *[
+                GRUCell(
+                    self.hidden_size,
+                    self.hidden_size,
+                    num_steps=self.num_steps,
+                    grad_slope=grad_slope,
+                    output_mems=(i == self.layers - 1),
+                )
+                for i in range(self.layers)
+            ]
+        )
+
+        self.__output_size = self.hidden_size 
+        self.fc = nn.Linear(self.__output_size, self.pre_length)
+
+        self.to('cuda:0') 
+
+    def forward(
+        self,
+        inputs: torch.Tensor,
+    ):
+        utils.reset(self.encoder)
+        for layer in self.net:
+            utils.reset(layer)
+
+
+        bs, length, c_num = inputs.size()
+
+        if self.args.normalize:
+            mean = inputs.mean(dim=1, keepdim=True).detach() # shape [B, 1, D]
+            inputs = inputs - mean
+            std = torch.sqrt(torch.var(inputs, dim=1, keepdim=True, unbiased=False) + 1e-5).detach()
+            inputs = inputs / std
+
+        h = self.encoder(inputs)  # B, H, C, L
+        hidden_size = h.size(1)
+        h = h.permute(0, 2, 3, 1).reshape(bs * c_num, length, hidden_size)  # BC, L, H
+        for i in range(length):
+            spks, mems = self.net(h[:, i, :])
+        spks = spks.reshape(bs, c_num * hidden_size, -1)  # B, CH, Time Step
+        spks = spks[:, :, -1]  # aggregate over time dimension shape, (B, CH)
+        preds = self.fc(spks.view(bs, c_num, -1)).squeeze(-1) # B, O, C
+        preds = preds.permute(0, 2, 1).contiguous()
+
+        if self.args.normalize:
+            preds = preds * std + mean  # denormalize
+
+        aux = {'gate_l0': torch.tensor(0.0, device=preds.device)} # palceholder
+
+        return preds, aux
+
+    @property
+    def output_size(self):
+        return self.__output_size

model/SpikeRNN_CPG.py

@@ -0,0 +1,489 @@
+from typing import Optional
+from pathlib import Path
+import torch
+from torch import nn
+from spikingjelly.activation_based import surrogate, neuron, functional
+import math
+import copy
+
+
+tau = 2.0 
+backend = "torch"
+detach_reset = True
+
+
+
+def generate_ones_and_minus_ones_matrix(rows, cols):
+    random_matrix = torch.randint(0, 2, (rows, cols))
+    binary_matrix = torch.where(
+        random_matrix == 0,
+        -1 * torch.ones_like(random_matrix),
+        torch.ones_like(random_matrix),
+    )
+    return binary_matrix.float()
+
+
+class RandomPE(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        pe_mode="concat",
+        num_pe_neuron=10,
+        neuron_pe_scale=1000.0,
+        dropout=0.1,
+        num_steps=4,
+    ):
+        super().__init__()
+        self.max_len = 5000  # different from windows
+        self.pe_mode = pe_mode
+        self.neuron_pe_scale = neuron_pe_scale
+        self.dropout = nn.Dropout(p=dropout)
+        if self.pe_mode == "concat":
+            self.num_pe_neuron = copy.deepcopy(num_pe_neuron)
+        elif self.pe_mode == "add":
+            self.num_pe_neuron = copy.deepcopy(d_model)
+        pe = generate_ones_and_minus_ones_matrix(
+            self.max_len, self.num_pe_neuron
+        )  # MaxL, Neur
+        pe = pe.unsqueeze(0).transpose(0, 1)  # MaxL, 1, Neur
+        print("pe.shape: ", pe.shape)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        # T, B, L, D
+        T, B, L, _ = x.shape
+        x = x.permute(1, 0, 2, 3)  # B, T, L, D
+        x = x.flatten(1, 2)  # B, TL, D
+        if self.pe_mode == "concat":
+            # tmp: TL, 1, Neur -> TL, B, Neur -> B, TL, Neur
+            tmp = self.pe[: x.size(-2), :].repeat(1, B, 1).transpose(0, 1)
+            x = torch.concat([x, tmp], dim=-1)
+            # print(x.shape) # B, TL, D'
+        elif self.pe_mode == "add":
+            # [B, TL, D] + [1, TL, Neur]
+            x = x + self.pe[: x.size(-2), :].transpose(0, 1)
+            # print(x.shape) # B, TL, D
+        x = x.transpose(0, 1)  # TL, B D
+        x = x.reshape(T, L, B, -1)  # T, L, B, D
+        x = x.permute(0, 2, 1, 3)  # T, B, L, D
+        return self.dropout(x)
+
+
+class NeuronPE(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        pe_mode="concat",
+        num_pe_neuron=10,
+        neuron_pe_scale=10000.0,
+        dropout=0.1,
+        num_steps=4,
+    ):
+        super().__init__()
+        self.max_len = 50000  # different from windows
+        self.pe_mode = pe_mode
+        self.neuron_pe_scale = neuron_pe_scale
+        self.dropout = nn.Dropout(p=dropout)
+        if self.pe_mode == "concat":
+            self.num_pe_neuron = copy.deepcopy(num_pe_neuron)
+        elif self.pe_mode == "add":
+            self.num_pe_neuron = copy.deepcopy(d_model)
+        pe = torch.zeros(self.max_len, self.num_pe_neuron)  # MaxL, Neur
+        position = torch.arange(0, self.max_len, dtype=torch.float).unsqueeze(
+            1
+        )  # MaxL, 1
+        div_term = torch.exp(
+            torch.arange(0, self.num_pe_neuron, 2).float()
+            * (-math.log(neuron_pe_scale) / self.num_pe_neuron)
+        )
+        div_term_single = torch.exp(
+            torch.arange(0, self.num_pe_neuron - 1, 2).float()
+            * (-math.log(neuron_pe_scale) / self.num_pe_neuron)
+        )
+        pe[:, 0::2] = torch.heaviside(
+            torch.sin(position * div_term) - 0.8, torch.tensor([1.0])
+        )
+        pe[:, 1::2] = torch.heaviside(
+            torch.cos(position * div_term_single) - 0.8, torch.tensor([1.0])
+        )
+        pe = pe.unsqueeze(0).transpose(0, 1)  # MaxL, 1, Neur
+        print("pe.shape: ", pe.shape)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        # T, B, L, D
+        T, B, L, _ = x.shape
+        x = x.permute(1, 0, 2, 3)  # B, T, L, D
+        x = x.flatten(1, 2)  # B, TL, D
+        if self.pe_mode == "concat":
+            # tmp: TL, 1, Neur -> TL, B, Neur -> B, TL, Neur
+            tmp = self.pe[: x.size(-2), :].repeat(1, B, 1).transpose(0, 1)
+            x = torch.concat([x, tmp], dim=-1)
+            # print(x.shape) # B, TL, D'
+        elif self.pe_mode == "add":
+            # [B, TL, D] + [1, TL, Neur]
+            x = x + self.pe[: x.size(-2), :].transpose(0, 1)
+            # print(x.shape) # B, TL, D
+        x = x.transpose(0, 1)  # TL, B D
+        x = x.reshape(T, L, B, -1)  # T, L, B, D
+        x = x.permute(0, 2, 1, 3)  # T, B, L, D
+        return self.dropout(x)
+
+
+class StaticPE(nn.Module):
+    r"""Inject some information about the relative or absolute position of the tokens
+        in the sequence. The positional encodings have the same dimension as
+        the embeddings, so that the two can be summed. Here, we use sine and cosine
+        functions of different frequencies.
+    .. math::
+        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
+        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
+        \text{where pos is the word position and i is the embed idx)"""
+
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        pe = torch.zeros(max_len, d_model)  # MaxL, D
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # MaxL, 1
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
+        )
+        div_term_single = torch.exp(
+            torch.arange(0, d_model - 1, 2).float() * (-math.log(10000.0) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term_single)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        # x: L, TB, D
+        x = x + self.pe[: x.size(0), :]
+        x = self.dropout(x)
+        return x
+
+
+class ConvPE(nn.Module):
+    def __init__(self, d_model, dropout=0.1, max_len=5000, num_steps=4):
+
+        super().__init__()
+        self.T = num_steps
+        self.rpe_conv = nn.Conv1d(
+            d_model, d_model, kernel_size=3, stride=1, padding=1, bias=False
+        )
+        self.rpe_bn = nn.BatchNorm1d(d_model)
+        self.rpe_lif = neuron.LIFNode(
+            step_mode="m",
+            detach_reset=True,
+            surrogate_function=surrogate.ATan(),
+            v_threshold=1.0,
+        )
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x):
+        # x: L, TB, D
+        L, TB, D = x.shape
+        x_feat = x.permute(1, 2, 0)  # TB, D, L
+        x_feat = self.rpe_conv(x_feat)  # TB, D, L
+        x_feat = (
+            self.rpe_bn(x_feat).reshape(self.T, int(TB / self.T), D, L).contiguous()
+        )  # T, B, D, L
+        x_feat = self.rpe_lif(x_feat)
+        x_feat = x_feat.flatten(0, 1)  # TB, D, L
+        x_feat = self.dropout(x_feat)  # TB, D, L
+        x_feat = x_feat.permute(2, 0, 1)  # L, TB, D
+        x = x + x_feat
+        return x
+
+
+class PositionEmbedding(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        pe_type: str,
+        max_len: int = 5000,
+        pe_mode: str = "add",
+        num_pe_neuron: int = 10,
+        neuron_pe_scale: float = 1000.0,
+        dropout=0.1,
+        num_steps=4,
+    ):
+        super().__init__()
+        self.emb_type = pe_type
+        if pe_type in ["learn", "none"]:
+            self.emb = nn.Embedding(max_len, input_size)
+        elif pe_type == "conv":
+            self.emb = ConvPE(
+                d_model=input_size,
+                max_len=max_len,
+                dropout=dropout,
+                num_steps=num_steps,
+            )
+        elif pe_type == "static":
+            self.emb = StaticPE(d_model=input_size, max_len=max_len, dropout=dropout)
+        elif pe_type == "neuron":
+            self.emb = NeuronPE(
+                d_model=input_size,
+                pe_mode=pe_mode,
+                num_pe_neuron=num_pe_neuron,
+                neuron_pe_scale=neuron_pe_scale,
+                dropout=dropout,
+                num_steps=num_steps,
+            )
+        elif pe_type == "random":
+            self.emb = RandomPE(
+                d_model=input_size,
+                pe_mode=pe_mode,
+                num_pe_neuron=num_pe_neuron,
+                neuron_pe_scale=neuron_pe_scale,
+                dropout=dropout,
+                num_steps=num_steps,
+            )
+        else:
+            raise ValueError("Unknown embedding type: {}".format(pe_type))
+
+    def forward(self, x):
+        if self.emb_type == "learn":
+            # T, B, L, D = x.shape # x: T, B, L, D
+            # x = x.flatten(0, 1) # TB, L, D
+            tmp = torch.arange(
+                end=x.size()[1], device=x.device
+            )  # [0,1,2,...,L-1], shape: L
+            embedding = self.emb(tmp)  # shape: L, D
+            embedding = embedding.repeat([x.size()[0], 1, 1])  # TB, L, D'
+            x = x + embedding
+            # x = x.reshape(T, B, L, -1)
+        elif self.emb_type in ["static", "conv"]:
+            T, B, L, _ = x.shape  # x: T, B, L, D
+            x = x.flatten(0, 1)  # TB, L, D
+            x = self.emb(x.transpose(0, 1)).transpose(0, 1)  # x: TB, L, D'
+            x = x.reshape(T, B, L, -1)
+        elif self.emb_type in ["neuron", "random"]:
+            T, B, L, _ = x.shape  # x: T, B, L, D
+            # T, B, L, D
+            x = self.emb(x)
+            x = x.reshape(T, B, L, -1)
+        return x  # T, B, L, D'
+
+
+class RepeatEncoder(nn.Module):
+    def __init__(self, output_size: int):
+        super().__init__()
+        self.out_size = output_size
+        self.lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: B, L, C
+        inputs = inputs.repeat(
+            tuple([self.out_size] + torch.ones(len(inputs.size()), dtype=int).tolist())
+        )  # T B L C
+        inputs = inputs.permute(0, 1, 3, 2)  # T B C L
+        spks = self.lif(inputs)  # T B C L
+        return spks
+
+
+class DeltaEncoder(nn.Module):
+    def __init__(self, output_size: int):
+        super().__init__()
+        self.norm = nn.BatchNorm2d(1)
+        self.enc = nn.Linear(1, output_size)
+        self.lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: B, L, C
+        delta = torch.zeros_like(inputs)
+        delta[:, 1:] = inputs[:, 1:, :] - inputs[:, :-1, :]
+        delta = delta.unsqueeze(1).permute(0, 1, 3, 2)  # B, 1, C, L
+        delta = self.norm(delta)
+        delta = delta.permute(0, 2, 3, 1)  # B, C, L, 1
+        enc = self.enc(delta)  # B, C, L, T
+        enc = enc.permute(3, 0, 1, 2)  # T, B, C, L
+        spks = self.lif(enc)
+        return spks
+
+
+class ConvEncoder(nn.Module):
+    def __init__(self, output_size: int, kernel_size: int = 3):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv2d(
+                in_channels=1,
+                out_channels=output_size,
+                kernel_size=(1, kernel_size),
+                stride=1,
+                padding=(0, kernel_size // 2),
+            ),
+            nn.BatchNorm2d(output_size),
+        )
+        self.lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: B, L, C
+        inputs = inputs.permute(0, 2, 1).unsqueeze(1)  # B, 1, C, L
+        enc = self.encoder(inputs)  # B, T, C, L
+        enc = enc.permute(1, 0, 2, 3)  # T, B, C, L
+        spks = self.lif(enc)  # T, B, C, L
+        return spks
+
+
+
+SpikeEncoder = {
+    "snntorch": {
+        "repeat": RepeatEncoder,
+        "conv": ConvEncoder,
+        "delta": DeltaEncoder,
+    },
+    "spikingjelly": {
+        "repeat": RepeatEncoder,
+        "conv": ConvEncoder,
+        "delta": DeltaEncoder,
+    },
+}
+
+
+
+class SpikeRNNCell(nn.Module):
+    def __init__(self, input_size: int, output_size: int):
+        super().__init__()
+        self.input_size = input_size
+        self.linear = nn.Linear(input_size, output_size)
+        self.lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+        )
+
+    def forward(self, x):
+        # T, B, L, C'
+        T, B, L, _ = x.shape
+        x = x.flatten(0, 1)  # TB, L, C'
+        x = self.linear(x)
+        x = x.reshape(T, B, L, -1)
+        x = self.lif(x)  # T, B, L, C'
+        return x
+
+
+class SpikeRNN_CPG(nn.Module):
+
+    def __init__(
+        self,
+        args,
+        hidden_size: int,
+        layers: int = 1,
+        num_steps: int = 4,
+        input_size: Optional[int] = None,
+        max_length: Optional[int] = 5000,
+        weight_file: Optional[Path] = None,
+        encoder_type: Optional[str] = "conv",
+        num_pe_neuron: int = 40,
+        pe_type: str = "neuron",
+        pe_mode: str = "concat",  # "add" or concat
+        neuron_pe_scale: float = 10000.0,  # "100" or "1000" or "10000"
+    ):
+        super().__init__()
+        self._snn_backend = "spikingjelly"
+        self.hidden_size   = args.hidden_size
+        self.num_steps   = args.T
+        self.input_size = args.feature_size
+        self.pre_length   = args.pre_length
+        self.layers       = args.blocks
+        self.pe_type = pe_type
+        self.pe_mode = pe_mode
+        self.num_pe_neuron = num_pe_neuron
+        self.neuron_pe_scale = neuron_pe_scale
+        self.temporal_encoder = SpikeEncoder[self._snn_backend][encoder_type](self.num_steps)
+        self.args = args
+
+        self.pe = PositionEmbedding(
+            pe_type=pe_type,
+            pe_mode=pe_mode,
+            neuron_pe_scale=neuron_pe_scale,
+            input_size=self.input_size,
+            max_len=max_length,
+            num_pe_neuron=self.num_pe_neuron,
+            dropout=0.1,
+            num_steps=self.num_steps,
+        )
+
+        if self.pe_type == "neuron" and self.pe_mode == "concat":
+            self.dim = hidden_size + num_pe_neuron
+        else:
+            self.dim = hidden_size
+
+        if self.pe_type == "neuron" and self.pe_mode == "concat":
+            self.encoder = nn.Linear(input_size + num_pe_neuron, self.dim)
+        else:
+            self.encoder = nn.Linear(input_size, self.dim)
+
+        self.init_lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+            v_threshold=1.0,
+            backend=backend,
+        )
+
+        self.net = nn.Sequential(
+            *[
+                SpikeRNNCell(input_size=self.dim, output_size=self.dim)
+                for i in range(layers)
+            ]
+        )
+
+        self.__output_size = self.dim
+        self.fc1 = nn.Linear(self.__output_size, args.feature_size)
+        self.fc2 = nn.Linear(args.seq_length, self.pre_length)
+        self.to('cuda:0') 
+
+
+    def forward(
+        self,
+        inputs: torch.Tensor,
+    ):
+        functional.reset_net(self)
+        if self.args.normalize:
+            mean = inputs.mean(dim=1, keepdim=True).detach() # shape [B, 1, D]
+            inputs = inputs - mean
+
+            std = torch.sqrt(torch.var(inputs, dim=1, keepdim=True, unbiased=False) + 1e-5).detach()
+            inputs = inputs / std
+
+
+        hiddens = self.temporal_encoder(inputs)  # T, B, C, L
+        hiddens = hiddens.transpose(-2, -1)  # T, B, L, C
+        T, B, L, _ = hiddens.size()  # T, B, L, D
+        if self.pe_type != "none":
+            hiddens = self.pe(hiddens)  # T B L C'
+        hiddens = self.encoder(hiddens.flatten(0, 1)).reshape(T, B, L, -1)  # T B L D
+        hiddens = self.init_lif(hiddens)
+        hiddens = self.net(hiddens)  # T, B, L, D
+        out = hiddens.mean(0) # B, L, D
+        preds = self.fc1(out)  # B, L, C
+        preds = self.fc2(preds.permute(0, 2, 1))  # B, C, L
+        preds = preds.permute(0, 2, 1).contiguous()
+
+        if self.args.normalize:
+            preds = preds * std + mean  # denormalize
+            
+        aux = {'gate_l0': torch.tensor(0.0, device=preds.device)} # placeholder
+        
+        return preds, aux
+
+

model/SpikeTCN_CPG.py

@@ -0,0 +1,596 @@
+from typing import Optional
+
+import torch
+from torch import nn
+from torch.nn.utils import weight_norm
+import snntorch as snn
+from snntorch import surrogate
+from snntorch import utils
+import copy
+import math
+
+def generate_ones_and_minus_ones_matrix(rows, cols):
+    random_matrix = torch.randint(0, 2, (rows, cols))
+    binary_matrix = torch.where(
+        random_matrix == 0,
+        -1 * torch.ones_like(random_matrix),
+        torch.ones_like(random_matrix),
+    )
+    return binary_matrix.float()
+
+
+class RandomPE(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        pe_mode="concat",
+        num_pe_neuron=10,
+        neuron_pe_scale=1000.0,
+        dropout=0.1,
+        num_steps=4,
+    ):
+        super().__init__()
+        self.max_len = 5000  # different from windows
+        self.pe_mode = pe_mode
+        self.neuron_pe_scale = neuron_pe_scale
+        self.dropout = nn.Dropout(p=dropout)
+        if self.pe_mode == "concat":
+            self.num_pe_neuron = copy.deepcopy(num_pe_neuron)
+        elif self.pe_mode == "add":
+            self.num_pe_neuron = copy.deepcopy(d_model)
+        pe = generate_ones_and_minus_ones_matrix(
+            self.max_len, self.num_pe_neuron
+        )  # MaxL, Neur
+        pe = pe.unsqueeze(0).transpose(0, 1)  # MaxL, 1, Neur
+        print("pe.shape: ", pe.shape)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        # T, B, L, D
+        T, B, L, _ = x.shape
+        x = x.permute(1, 0, 2, 3)  # B, T, L, D
+        x = x.flatten(1, 2)  # B, TL, D
+        if self.pe_mode == "concat":
+            # tmp: TL, 1, Neur -> TL, B, Neur -> B, TL, Neur
+            tmp = self.pe[: x.size(-2), :].repeat(1, B, 1).transpose(0, 1)
+            x = torch.concat([x, tmp], dim=-1)
+            # print(x.shape) # B, TL, D'
+        elif self.pe_mode == "add":
+            # [B, TL, D] + [1, TL, Neur]
+            x = x + self.pe[: x.size(-2), :].transpose(0, 1)
+            # print(x.shape) # B, TL, D
+        x = x.transpose(0, 1)  # TL, B D
+        x = x.reshape(T, L, B, -1)  # T, L, B, D
+        x = x.permute(0, 2, 1, 3)  # T, B, L, D
+        return self.dropout(x)
+
+
+class NeuronPE(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        pe_mode="concat",
+        num_pe_neuron=10,
+        neuron_pe_scale=10000.0,
+        dropout=0.1,
+        num_steps=4,
+    ):
+        super().__init__()
+        self.max_len = 50000  # different from windows
+        self.pe_mode = pe_mode
+        self.neuron_pe_scale = neuron_pe_scale
+        self.dropout = nn.Dropout(p=dropout)
+        if self.pe_mode == "concat":
+            self.num_pe_neuron = copy.deepcopy(num_pe_neuron)
+        elif self.pe_mode == "add":
+            self.num_pe_neuron = copy.deepcopy(d_model)
+        pe = torch.zeros(self.max_len, self.num_pe_neuron)  # MaxL, Neur
+        position = torch.arange(0, self.max_len, dtype=torch.float).unsqueeze(
+            1
+        )  # MaxL, 1
+        div_term = torch.exp(
+            torch.arange(0, self.num_pe_neuron, 2).float()
+            * (-math.log(neuron_pe_scale) / self.num_pe_neuron)
+        )
+        div_term_single = torch.exp(
+            torch.arange(0, self.num_pe_neuron - 1, 2).float()
+            * (-math.log(neuron_pe_scale) / self.num_pe_neuron)
+        )
+        pe[:, 0::2] = torch.heaviside(
+            torch.sin(position * div_term) - 0.8, torch.tensor([1.0])
+        )
+        pe[:, 1::2] = torch.heaviside(
+            torch.cos(position * div_term_single) - 0.8, torch.tensor([1.0])
+        )
+        pe = pe.unsqueeze(0).transpose(0, 1)  # MaxL, 1, Neur
+        print("pe.shape: ", pe.shape)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        # T, B, L, D
+        T, B, L, _ = x.shape
+        x = x.permute(1, 0, 2, 3)  # B, T, L, D
+        x = x.flatten(1, 2)  # B, TL, D
+        if self.pe_mode == "concat":
+            # tmp: TL, 1, Neur -> TL, B, Neur -> B, TL, Neur
+            tmp = self.pe[: x.size(-2), :].repeat(1, B, 1).transpose(0, 1)
+            x = torch.concat([x, tmp], dim=-1)
+            # print(x.shape) # B, TL, D'
+        elif self.pe_mode == "add":
+            # [B, TL, D] + [1, TL, Neur]
+            # print(self.pe[:x.size(-2), :].shape)
+            x = x + self.pe[: x.size(-2), :].transpose(0, 1)
+            # print(x.shape) # B, TL, D
+        x = x.transpose(0, 1)  # TL, B D
+        x = x.reshape(T, L, B, -1)  # T, L, B, D
+        x = x.permute(0, 2, 1, 3)  # T, B, L, D
+        return self.dropout(x)
+
+
+class StaticPE(nn.Module):
+    r"""Inject some information about the relative or absolute position of the tokens
+        in the sequence. The positional encodings have the same dimension as
+        the embeddings, so that the two can be summed. Here, we use sine and cosine
+        functions of different frequencies.
+    .. math::
+        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
+        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
+        \text{where pos is the word position and i is the embed idx)"""
+
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        pe = torch.zeros(max_len, d_model)  # MaxL, D
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # MaxL, 1
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
+        )
+        div_term_single = torch.exp(
+            torch.arange(0, d_model - 1, 2).float() * (-math.log(10000.0) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term_single)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        # x: L, TB, D
+        x = x + self.pe[: x.size(0), :]
+        x = self.dropout(x)
+        return x
+
+
+class ConvPE(nn.Module):
+    def __init__(self, d_model, dropout=0.1, max_len=5000, num_steps=4):
+
+        super().__init__()
+        self.T = num_steps
+        self.rpe_conv = nn.Conv1d(
+            d_model, d_model, kernel_size=3, stride=1, padding=1, bias=False
+        )
+        self.rpe_bn = nn.BatchNorm1d(d_model)
+        self.rpe_lif = neuron.LIFNode(
+            step_mode="m",
+            detach_reset=True,
+            surrogate_function=surrogate.ATan(),
+            v_threshold=1.0,
+        )
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x):
+        # x: L, TB, D
+        L, TB, D = x.shape
+        x_feat = x.permute(1, 2, 0)  # TB, D, L
+        x_feat = self.rpe_conv(x_feat)  # TB, D, L
+        x_feat = (
+            self.rpe_bn(x_feat).reshape(self.T, int(TB / self.T), D, L).contiguous()
+        )  # T, B, D, L
+        x_feat = self.rpe_lif(x_feat)
+        x_feat = x_feat.flatten(0, 1)  # TB, D, L
+        x_feat = self.dropout(x_feat)  # TB, D, L
+        x_feat = x_feat.permute(2, 0, 1)  # L, TB, D
+        x = x + x_feat
+        return x
+
+
+class PositionEmbedding(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        pe_type: str,
+        max_len: int = 5000,
+        pe_mode: str = "add",
+        num_pe_neuron: int = 10,
+        neuron_pe_scale: float = 1000.0,
+        dropout=0.1,
+        num_steps=4,
+    ):
+        super().__init__()
+        self.emb_type = pe_type
+        if pe_type in ["learn", "none"]:
+            self.emb = nn.Embedding(max_len, input_size)
+        elif pe_type == "conv":
+            self.emb = ConvPE(
+                d_model=input_size,
+                max_len=max_len,
+                dropout=dropout,
+                num_steps=num_steps,
+            )
+        elif pe_type == "static":
+            self.emb = StaticPE(d_model=input_size, max_len=max_len, dropout=dropout)
+        elif pe_type == "neuron":
+            self.emb = NeuronPE(
+                d_model=input_size,
+                pe_mode=pe_mode,
+                num_pe_neuron=num_pe_neuron,
+                neuron_pe_scale=neuron_pe_scale,
+                dropout=dropout,
+                num_steps=num_steps,
+            )
+        elif pe_type == "random":
+            self.emb = RandomPE(
+                d_model=input_size,
+                pe_mode=pe_mode,
+                num_pe_neuron=num_pe_neuron,
+                neuron_pe_scale=neuron_pe_scale,
+                dropout=dropout,
+                num_steps=num_steps,
+            )
+        else:
+            raise ValueError("Unknown embedding type: {}".format(pe_type))
+
+    def forward(self, x):
+        if self.emb_type == "learn":
+            # T, B, L, D = x.shape # x: T, B, L, D
+            # x = x.flatten(0, 1) # TB, L, D
+            tmp = torch.arange(
+                end=x.size()[1], device=x.device
+            )  # [0,1,2,...,L-1], shape: L
+            embedding = self.emb(tmp)  # shape: L, D
+            embedding = embedding.repeat([x.size()[0], 1, 1])  # TB, L, D'
+            x = x + embedding
+            # x = x.reshape(T, B, L, -1)
+        elif self.emb_type in ["static", "conv"]:
+            T, B, L, _ = x.shape  # x: T, B, L, D
+            x = x.flatten(0, 1)  # TB, L, D
+            x = self.emb(x.transpose(0, 1)).transpose(0, 1)  # x: TB, L, D'
+            x = x.reshape(T, B, L, -1)
+        elif self.emb_type in ["neuron", "random"]:
+            T, B, L, _ = x.shape  # x: T, B, L, D
+            # T, B, L, D
+            x = self.emb(x)
+            x = x.reshape(T, B, L, -1)
+        return x  # T, B, L, D'
+
+
+
+
+
+
+class RepeatEncoder(nn.Module):
+    def __init__(self, output_size: int):
+        super().__init__()
+        self.out_size = output_size
+        self.lif = snn.Leaky(
+            beta=0.99, spike_grad=surrogate.atan(), init_hidden=True, output=False
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: batch, L, C
+        inputs = inputs.repeat(
+            tuple([self.out_size] + torch.ones(len(inputs.size()), dtype=int).tolist())
+        )  # out_size batch L C
+        inputs = inputs.permute(1, 0, 3, 2)  # batch out_size L C
+        spks = self.lif(inputs)
+        return spks
+
+
+class ConvEncoder(nn.Module):
+    def __init__(self, output_size: int, kernel_size: int = 3):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv2d(
+                in_channels=1,
+                out_channels=output_size,
+                kernel_size=(1, kernel_size),
+                stride=1,
+                padding=(0, kernel_size // 2),
+            ),
+            nn.BatchNorm2d(output_size),
+        )
+        self.lif = snn.Leaky(
+            beta=0.99,
+            spike_grad=surrogate.atan(alpha=2.0),
+            init_hidden=True,
+            output=False,
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: batch, L, C
+        inputs = inputs.permute(0, 2, 1).unsqueeze(1)  # batch, 1, C, L
+        enc = self.encoder(inputs)  # batch, output_size, C, L
+        spks = self.lif(enc)
+        return spks
+
+
+class DeltaEncoder(nn.Module):
+    def __init__(self, output_size: int):
+        super().__init__()
+        self.norm = nn.BatchNorm2d(1)
+        self.enc = nn.Linear(1, output_size)
+        self.lif = snn.Leaky(
+            beta=0.99, spike_grad=surrogate.atan(), init_hidden=True, output=False
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: batch, L, C
+        delta = torch.zeros_like(inputs)
+        delta[:, 1:] = inputs[:, 1:, :] - inputs[:, :-1, :]
+        delta = delta.unsqueeze(1).permute(0, 1, 3, 2)  # batch, 1, C, L
+        delta = self.norm(delta)
+        delta = delta.permute(0, 2, 3, 1)  # batch, C, L, 1
+        enc = self.enc(delta)  # batch, C, L, output_size
+        enc = enc.permute(0, 3, 1, 2)  # batch, output_size, C, L
+        spks = self.lif(enc)
+        return spks
+
+
+
+class Chomp1d(nn.Module):
+    def __init__(self, chomp_size):
+        super().__init__()
+        self.chomp_size = chomp_size
+
+    def forward(self, x):
+        return x[:, :, : -self.chomp_size].contiguous()
+
+
+class Chomp2d(nn.Module):
+    def __init__(self, chomp_size):
+        super().__init__()
+        self.chomp_size = chomp_size
+
+    def forward(self, x):
+        return x[:, :, :, : -self.chomp_size].contiguous()
+
+
+
+SpikeEncoder = {
+    "snntorch": {
+        "repeat": RepeatEncoder,
+        "conv": ConvEncoder,
+        "delta": DeltaEncoder,
+    },
+    "spikingjelly": {
+        "repeat": RepeatEncoder,
+        "conv": ConvEncoder,
+        "delta": DeltaEncoder,
+    },
+}
+
+
+
+class SpikeTemporalBlock2D(nn.Module):
+    def __init__(
+        self,
+        n_inputs,
+        n_outputs,
+        kernel_size,
+        stride,
+        dilation,
+        padding,
+        num_steps=4,
+    ):
+        super().__init__()
+        self.num_steps = num_steps
+        self.conv1 = weight_norm(
+            nn.Conv2d(
+                n_inputs,
+                n_outputs,
+                (1, kernel_size),
+                stride=stride,
+                padding=(0, padding),
+                dilation=(1, dilation),
+            )
+        )
+        self.bn1 = nn.BatchNorm2d(n_outputs)
+        self.chomp1 = Chomp2d(padding)
+        self.lif1 = snn.Leaky(
+            beta=0.99,
+            spike_grad=surrogate.atan(alpha=2.0),
+            init_hidden=True,
+            threshold=1.0,
+        )
+
+        self.conv2 = weight_norm(
+            nn.Conv2d(
+                n_outputs,
+                n_outputs,
+                (1, kernel_size),
+                stride=stride,
+                padding=(0, padding),
+                dilation=(1, dilation),
+            )
+        )
+        self.bn2 = nn.BatchNorm2d(n_outputs)
+        self.chomp2 = Chomp2d(padding)
+        self.lif2 = snn.Leaky(
+            beta=0.99,
+            spike_grad=surrogate.atan(alpha=2.0),
+            init_hidden=True,
+            threshold=1.0,
+        )
+
+        self.downsample = (
+            nn.Conv2d(n_inputs, n_outputs, (1, 1)) if n_inputs != n_outputs else None
+        )
+        self.lif = snn.Leaky(
+            beta=0.99,
+            spike_grad=surrogate.atan(alpha=2.0),
+            init_hidden=True,
+            threshold=1.0,
+        )
+
+    def init_weights(self):
+        self.conv1.weight.data.normal_(0, 0.01)
+        self.conv2.weight.data.normal_(0, 0.01)
+        if self.downsample is not None:
+            self.downsample.weight.data.normal_(0, 0.01)
+
+    def forward(self, x):
+        out1 = self.chomp1(self.bn1(self.conv1(x)))
+        spk_rec1 = []
+        for _ in range(self.num_steps):
+            spk = self.lif1(out1)
+            spk_rec1.append(spk)
+        spks1 = torch.stack(spk_rec1, dim=-1)  # spks1: B, H, C, L, T
+        spks1 = spks1.mean(-1)  # spks1: B, H, C, L
+
+        out2 = self.chomp2(self.bn2(self.conv2(spks1)))
+        spk_rec2 = []
+        for _ in range(self.num_steps):
+            spk = self.lif2(out2)
+            spk_rec2.append(spk)
+        spks2 = torch.stack(spk_rec2, dim=-1)  # spks2: B, H, C, L, T
+        spks2 = spks2.mean(-1)  # spks2: B, H, C, L
+
+        if torch.isnan(spks2).any() or torch.isinf(spks2).any():
+            print("illegal value in TemporalBlock2D")
+
+        if self.downsample is None:
+            res = x
+        else:
+            res = self.downsample(x)
+        spk_rec3 = []
+        for _ in range(self.num_steps):
+            spk = self.lif(spks2 + res)
+            spk_rec3.append(spk)
+
+        res = torch.stack(spk_rec3, dim=-1)  # res: B, H, C, L, T
+        res = res.mean(-1)
+
+        return res
+
+
+class SpikeTCN_CPG(nn.Module):
+    
+
+    def __init__(
+        self,
+        args,
+        num_levels: int=3,
+        channel: int=16,
+        dilation: int=2,
+        stride: int = 1,
+        num_steps: int = 16,
+        kernel_size: int = 2,
+        dropout: float = 0.2,
+        max_length: int = 100,
+        input_size: Optional[int] = None,
+        hidden_size: int = 128,
+        encoder_type: Optional[str] = "conv",
+        num_pe_neuron: int = 40,
+        pe_type: str = "neuron",
+        pe_mode: str = "concat",  # "add" or "concat"
+        neuron_pe_scale: float = 10000.0,  # "100" or "1000" or "10000"
+    ):
+        """
+        Args:
+            num_channels: The number of convolutional channels in each layer.
+            kernel_size: The kernel size of convolutional layers.
+            dropout: Dropout rate.
+        """
+        super().__init__()
+        self.pe_type = pe_type
+        self._snn_backend = "snntorch"
+        self.pe_mode = pe_mode
+        self.num_pe_neuron = num_pe_neuron
+        self.hidden_size   = args.hidden_size  
+        self.num_steps = args.T    
+        self.input_size = args.feature_size  
+        self.pre_length = args.pre_length      
+        self.num_levels = args.blocks
+        self.pe_type = pe_type
+        self.pe_mode = pe_mode
+        self.num_pe_neuron = num_pe_neuron
+        self.kernel_size = args.kernel_size
+
+        self.encoder = SpikeEncoder[self._snn_backend][encoder_type](self.hidden_size)
+        self.args = args
+
+ 
+        self.pe = PositionEmbedding(
+            pe_type=pe_type,
+            pe_mode=pe_mode,
+            neuron_pe_scale=neuron_pe_scale,
+            input_size=self.input_size,
+            max_len=max_length,
+            num_pe_neuron=self.num_pe_neuron,
+            dropout=0.1,
+            num_steps=self.num_steps,
+        )
+        layers = []
+        num_channels = [channel] * self.num_levels
+        num_channels.append(1)
+        for i in range(self.num_levels + 1):
+            dilation_size = dilation**i
+            in_channels = self.hidden_size if i == 0 else num_channels[i - 1]
+            out_channels = num_channels[i]
+            layers += [
+                SpikeTemporalBlock2D(
+                    in_channels,
+                    out_channels,
+                    self.kernel_size,
+                    stride=stride,
+                    dilation=dilation_size,
+                    padding=(self.kernel_size - 1) * dilation_size,
+                    num_steps=self.num_steps,
+                )
+            ]
+
+        self.network = nn.Sequential(*layers)
+        if (self.pe_type == "neuron" and self.pe_mode == "concat") or (
+            self.pe_type == "random" and self.pe_mode == "concat"
+        ):
+            self.__output_size = args.feature_size + num_pe_neuron
+        else:
+            self.__output_size =  args.seq_length
+
+        self.fc1 = nn.Linear(self.__output_size, args.feature_size)
+        self.fc2 = nn.Linear(args.seq_length, self.pre_length)
+        self.to('cuda:0') 
+
+    def forward(self, inputs: torch.Tensor):
+        utils.reset(self.encoder)
+        for layer in self.network:
+            utils.reset(layer)
+        
+        if self.args.normalize:
+            mean = inputs.mean(dim=1, keepdim=True).detach() # shape [B, 1, D]
+            inputs = inputs - mean
+
+            std = torch.sqrt(torch.var(inputs, dim=1, keepdim=True, unbiased=False) + 1e-5).detach()
+            inputs = inputs / std
+
+        inputs = self.encoder(inputs)  # B, H, C, L
+        if self.pe_type != "none":
+            # B, H, C, L -> H B L C' -> B H C' L
+            inputs = self.pe(inputs.permute(1, 0, 3, 2)).permute(1, 0, 3, 2)
+        spks = self.network(inputs)
+        spks = spks.squeeze(1)  # B, C', L
+
+        preds = self.fc1(spks.permute(0, 2, 1))  # B, L, C
+        preds = self.fc2(preds.permute(0, 2, 1))  # B, C', L
+        #.squeeze(-1) # B, O, C'
+        preds = preds.permute(0, 2, 1).contiguous()
+        if self.args.normalize:
+            preds = preds * std + mean  # denormalize
+        aux = {'gate_l0': torch.tensor(0.0, device=preds.device)} # placeholder
+        
+        return preds, aux
+
+    @property
+    def output_size(self):
+        return self.__output_size
+
+

model/Spikformer_CPG.py

@@ -0,0 +1,487 @@
+from typing import Optional
+
+from pathlib import Path
+import torch
+from torch import nn
+from spikingjelly.activation_based import surrogate, neuron, functional
+
+import math
+from dataclasses import dataclass
+import warnings
+
+
+
+tau = 2.0  # beta = 1 - 1/tau
+backend = "torch"
+detach_reset = True
+
+
+
+@dataclass
+class CPG(nn.Module):
+    num_neurons: int = 40
+    w_max: float = 10000.0
+    l_max: int = 5000
+
+    def __post_init__(self):
+        self._cpg = torch.zeros(self.l_max, self.num_neurons)
+        position = torch.arange(0, self.l_max, dtype=torch.float).unsqueeze(
+            1
+        )  # MaxL, 1
+        div_term = torch.exp(
+            torch.arange(0, self.num_neurons, 2).float()
+            * (-math.log(self.w_max) / self.num_neurons)
+        )
+        div_term_single = torch.exp(
+            torch.arange(0, self.num_neurons - 1, 2).float()
+            * (-math.log(self.w_max) / self.num_neurons)
+        )
+        self._cpg[:, 0::2] = torch.heaviside(
+            torch.sin(position * div_term) - 0.8, torch.tensor([1.0])
+        )
+        self._cpg[:, 1::2] = torch.heaviside(
+            torch.cos(position * div_term_single) - 0.8, torch.tensor([1.0])
+        )
+
+    @property
+    def cpg(self):
+        return self._cpg
+
+
+class CPGLinear(nn.Module):
+    def __init__(
+        self, input_size: int, output_size: int, cpg: CPG = CPG(), dropout: float = 0.1
+    ):
+        super().__init__()
+        self.cpg = nn.Parameter(cpg.cpg, requires_grad=False)
+        self.inp_linear = nn.Linear(input_size, output_size)
+        self.cpg_linear = nn.Linear(cpg.num_neurons, output_size)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor):
+        # B TL D
+        cpg = self.cpg[: x.size(-2)]
+        x = self.dropout(x)
+        return self.inp_linear(x) + self.cpg_linear(cpg)
+
+
+
+
+class RepeatEncoder(nn.Module):
+    def __init__(self, output_size: int):
+        super().__init__()
+        self.out_size = output_size
+        self.lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: B, L, C
+        inputs = inputs.repeat(
+            tuple([self.out_size] + torch.ones(len(inputs.size()), dtype=int).tolist())
+        )  # T B L C
+        inputs = inputs.permute(0, 1, 3, 2)  # T B C L
+        spks = self.lif(inputs)  # T B C L
+        return spks
+
+
+class DeltaEncoder(nn.Module):
+    def __init__(self, output_size: int):
+        super().__init__()
+        self.norm = nn.BatchNorm2d(1)
+        self.enc = nn.Linear(1, output_size)
+        self.lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: B, L, C
+        delta = torch.zeros_like(inputs)
+        delta[:, 1:] = inputs[:, 1:, :] - inputs[:, :-1, :]
+        delta = delta.unsqueeze(1).permute(0, 1, 3, 2)  # B, 1, C, L
+        delta = self.norm(delta)
+        delta = delta.permute(0, 2, 3, 1)  # B, C, L, 1
+        enc = self.enc(delta)  # B, C, L, T
+        enc = enc.permute(3, 0, 1, 2)  # T, B, C, L
+        spks = self.lif(enc)
+        return spks
+
+
+class ConvEncoder(nn.Module):
+    def __init__(self, output_size: int, kernel_size: int = 3):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv2d(
+                in_channels=1,
+                out_channels=output_size,
+                kernel_size=(1, kernel_size),
+                stride=1,
+                padding=(0, kernel_size // 2),
+            ),
+            nn.BatchNorm2d(output_size),
+        )
+        self.lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: B, L, C
+        inputs = inputs.permute(0, 2, 1).unsqueeze(1)  # B, 1, C, L
+        enc = self.encoder(inputs)  # B, T, C, L
+        enc = enc.permute(1, 0, 2, 3)  # T, B, C, L
+        spks = self.lif(enc)  # T, B, C, L
+        return spks
+
+
+
+SpikeEncoder = {
+    "snntorch": {
+        "repeat": RepeatEncoder,
+        "conv": ConvEncoder,
+        "delta": DeltaEncoder,
+    },
+    "spikingjelly": {
+        "repeat": RepeatEncoder,
+        "conv": ConvEncoder,
+        "delta": DeltaEncoder,
+    },
+}
+
+
+
+
+
+class SSA(nn.Module):
+    def __init__(
+        self, length, tau, common_thr, dim, heads=8, qkv_bias=False, qk_scale=0.25
+    ):
+        super().__init__()
+        assert dim % heads == 0, f"dim {dim} should be divided by num_heads {heads}."
+
+        self.dim = dim
+        self.heads = heads
+        self.qk_scale = qk_scale
+
+        self.q_m = nn.Linear(dim, dim)
+        self.q_bn = nn.BatchNorm1d(dim)
+        self.q_lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+            v_threshold=common_thr,
+            backend=backend,
+        )
+
+        self.k_m = nn.Linear(dim, dim)
+        self.k_bn = nn.BatchNorm1d(dim)
+        self.k_lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+            v_threshold=common_thr,
+            backend=backend,
+        )
+
+        self.v_m = nn.Linear(dim, dim)
+        self.v_bn = nn.BatchNorm1d(dim)
+        self.v_lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+            v_threshold=common_thr,
+            backend=backend,
+        )
+
+        self.attn_lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+            v_threshold=common_thr / 2,
+            backend=backend,
+        )
+
+        self.last_m = nn.Linear(dim, dim)
+        self.last_bn = nn.BatchNorm1d(dim)
+        self.last_lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+            v_threshold=common_thr,
+            backend=backend,
+        )
+
+    def forward(self, x):
+        T, B, L, D = x.shape
+        x_for_qkv = x.flatten(0, 1)  # TB L D
+        q_m_out = self.q_m(x_for_qkv)  # TB L D
+        q_m_out = (
+            self.q_bn(q_m_out.transpose(-1, -2))
+            .transpose(-1, -2)
+            .reshape(T, B, L, D)
+            .contiguous()
+        )
+        q_m_out = self.q_lif(q_m_out)
+        q = (
+            q_m_out.reshape(T, B, L, self.heads, D // self.heads)
+            .permute(0, 1, 3, 2, 4)
+            .contiguous()
+        )
+
+        k_m_out = self.k_m(x_for_qkv)
+        k_m_out = (
+            self.k_bn(k_m_out.transpose(-1, -2))
+            .transpose(-1, -2)
+            .reshape(T, B, L, D)
+            .contiguous()
+        )
+        k_m_out = self.k_lif(k_m_out)
+        k = (
+            k_m_out.reshape(T, B, L, self.heads, D // self.heads)
+            .permute(0, 1, 3, 2, 4)
+            .contiguous()
+        )
+
+        v_m_out = self.v_m(x_for_qkv)
+        v_m_out = (
+            self.v_bn(v_m_out.transpose(-1, -2))
+            .transpose(-1, -2)
+            .reshape(T, B, L, D)
+            .contiguous()
+        )
+        v_m_out = self.v_lif(v_m_out)
+        v = (
+            v_m_out.reshape(T, B, L, self.heads, D // self.heads)
+            .permute(0, 1, 3, 2, 4)
+            .contiguous()
+        )
+
+        attn = (q @ k.transpose(-2, -1)) * self.qk_scale
+        x = attn @ v  # x_shape: T * B * heads * L * D//heads
+
+        x = x.transpose(2, 3).reshape(T, B, L, D).contiguous()
+        x = self.attn_lif(x)
+
+        x = x.flatten(0, 1)
+        x = self.last_m(x)
+        x = self.last_bn(x.transpose(-1, -2)).transpose(-1, -2)
+        x = self.last_lif(x.reshape(T, B, L, D).contiguous())
+        return x
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        length,
+        tau,
+        common_thr,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        self.in_features = in_features
+        self.hidden_features = hidden_features
+        self.out_features = out_features
+
+        self.fc1 = CPGLinear(in_features, hidden_features)
+        self.bn1 = nn.BatchNorm1d(hidden_features)
+        self.lif1 = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+            v_threshold=common_thr,
+            backend=backend,
+        )
+
+        self.fc2 = CPGLinear(hidden_features, out_features)
+        self.bn2 = nn.BatchNorm1d(out_features)
+        self.lif2 = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+            v_threshold=common_thr,
+            backend=backend,
+        )
+
+    def forward(self, x):
+        T, B, L, D = x.shape
+        x = x.transpose(0, 1).flatten(1, 2)  # B TL D
+        x = self.fc1(x)  # B TL H
+        x = (
+            self.bn1(x.transpose(-1, -2))
+            .transpose(-1, -2)
+            .reshape(B, T, L, self.hidden_features)
+            .contiguous()
+        )  # B T L H
+        x = self.lif1(x.transpose(0, 1)).transpose(0, 1)  # B T L H
+        x = x.flatten(1, 2)  # B TL H
+        x = self.fc2(x)  # B TL D
+        x = (
+            self.bn2(x.transpose(-1, -2))
+            .transpose(-1, -2)
+            .reshape(B, T, L, D)
+            .contiguous()
+        )  # B T L D
+        x = self.lif2(x.transpose(0, 1))  # T B L D
+        return x
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        length,
+        tau,
+        common_thr,
+        dim,
+        d_ff,
+        heads=8,
+        qkv_bias=False,
+        qk_scale=0.125,
+    ):
+        super().__init__()
+        self.attn = SSA(
+            length=length,
+            tau=tau,
+            common_thr=common_thr,
+            dim=dim,
+            heads=heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+        )
+        self.mlp = MLP(
+            length=length,
+            tau=tau,
+            common_thr=common_thr,
+            in_features=dim,
+            hidden_features=d_ff,
+        )
+
+    def forward(self, x):
+        # T B L D
+        x = x + self.attn(x)
+        x = x + self.mlp(x)
+        return x
+
+
+class Spikformer_CPG(nn.Module):
+    def __init__(
+        self,
+        args,
+        dim: int=256,
+        d_ff: Optional[int] = None,
+        num_pe_neuron: int = 40,
+        pe_type: str = "neuron",
+        pe_mode: str = "concat",  # "add" or concat
+        neuron_pe_scale: float = 10000.0,  # "100" or "1000" or "10000"
+        depths: int = 2,
+        common_thr: float = 1.0,
+        max_length: int = 5000,
+        num_steps: int = 4,
+        heads: int = 8,
+        qkv_bias: bool = False,
+        qk_scale: float = 0.125,
+        input_size: Optional[int] = None,
+        weight_file: Optional[Path] = None,
+    ):
+        super().__init__()
+        self.dim = 256
+        self.d_ff = 1024
+        self.T = args.T
+        self.depths = args.blocks
+        self.pe_type = pe_type
+        self.pe_mode = pe_mode
+        self.num_pe_neuron = num_pe_neuron
+        self.input_size = args.feature_size
+        self.pre_length = args.pre_length
+        self.args = args
+
+
+        self._snn_backend = "spikingjelly"
+
+        self.temporal_encoder = SpikeEncoder[self._snn_backend]["conv"](num_steps)
+        self.encoder = CPGLinear(self.input_size, dim, CPG(num_neurons=num_pe_neuron))
+
+        self.init_lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+            v_threshold=common_thr,
+            backend=backend,
+        )
+
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    length=max_length,
+                    tau=tau,
+                    common_thr=common_thr,
+                    dim=dim,
+                    d_ff=self.d_ff,
+                    heads=heads,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                )
+                for _ in range(depths)
+            ]
+        )
+
+        self.apply(self._init_weights)
+
+        self.fc = nn.Linear(args.seq_length*dim, args.pre_length*args.feature_size)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0.0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.weight, 1.0)
+            nn.init.constant_(m.bias, 0.0)
+
+    def forward(self, x: torch.Tensor):
+        functional.reset_net(self)
+
+        if self.args.normalize:
+
+            mean = x.mean(dim=1, keepdim=True).detach() # shape [B, 1, D]
+            x = x - mean
+
+            std = torch.sqrt(torch.var(x, dim=1, keepdim=True, unbiased=False) + 1e-5).detach()
+            x = x / std
+
+
+        x = self.temporal_encoder(x)  # B L C -> T B C L
+        T, B, _, L = x.shape
+        x = x.permute(1, 0, 3, 2)  # B T L C
+        x = x.flatten(1, 2)  # B TL C
+        x = self.encoder(x)  # B TL D
+        x = x.reshape(B, T, L, -1).permute(1, 0, 2, 3)  # T B L D
+        x = self.init_lif(x)
+
+        for blk in self.blocks:
+            x = blk(x)  # T B L D
+        out = x.mean(0)
+        out = self.fc(out.flatten(-2, -1)).reshape(-1, self.pre_length, self.input_size)  # B D L -> B L D
+        if self.args.normalize:
+            out = out * std + mean  # denormalization
+        aux = {'gate_l0': torch.tensor(0.0, device=out.device)} # placeholder
+        return out, aux  # B D L -> B L D
+

model/TS_Former.py

@@ -0,0 +1,1365 @@
+from typing import Optional, Callable
+
+from pathlib import Path
+import torch
+from torch import nn
+from spikingjelly.activation_based import surrogate, neuron, functional
+
+import math
+import copy
+from spikingjelly.activation_based import surrogate, neuron
+from abc import abstractmethod
+import snntorch as snn
+from snntorch import utils
+import warnings
+
+surrogate.ATan = lambda alpha=2.0: SG.apply
+
+
+def generate_ones_and_minus_ones_matrix(rows, cols):
+    random_matrix = torch.randint(0, 2, (rows, cols))
+    binary_matrix = torch.where(
+        random_matrix == 0,
+        -1 * torch.ones_like(random_matrix),
+        torch.ones_like(random_matrix),
+    )
+    return binary_matrix.float()
+
+
+class RandomPE(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        pe_mode="concat",
+        num_pe_neuron=10,
+        neuron_pe_scale=1000.0,
+        dropout=0.1,
+        num_steps=4,
+    ):
+        super().__init__()
+        self.max_len = 5000  # different from windows
+        self.pe_mode = pe_mode
+        self.neuron_pe_scale = neuron_pe_scale
+        self.dropout = nn.Dropout(p=dropout)
+        if self.pe_mode == "concat":
+            self.num_pe_neuron = copy.deepcopy(num_pe_neuron)
+        elif self.pe_mode == "add":
+            self.num_pe_neuron = copy.deepcopy(d_model)
+        pe = generate_ones_and_minus_ones_matrix(
+            self.max_len, self.num_pe_neuron
+        )  # MaxL, Neur
+        pe = pe.unsqueeze(0).transpose(0, 1)  # MaxL, 1, Neur
+        print("pe.shape: ", pe.shape)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        # T, B, L, D
+        T, B, L, _ = x.shape
+        x = x.permute(1, 0, 2, 3)  # B, T, L, D
+        x = x.flatten(1, 2)  # B, TL, D
+        if self.pe_mode == "concat":
+            # tmp: TL, 1, Neur -> TL, B, Neur -> B, TL, Neur
+            tmp = self.pe[: x.size(-2), :].repeat(1, B, 1).transpose(0, 1)
+            x = torch.concat([x, tmp], dim=-1)
+            # print(x.shape) # B, TL, D'
+        elif self.pe_mode == "add":
+            # [B, TL, D] + [1, TL, Neur]
+            x = x + self.pe[: x.size(-2), :].transpose(0, 1)
+            # print(x.shape) # B, TL, D
+        x = x.transpose(0, 1)  # TL, B D
+        x = x.reshape(T, L, B, -1)  # T, L, B, D
+        x = x.permute(0, 2, 1, 3)  # T, B, L, D
+        return self.dropout(x)
+
+
+class NeuronPE(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        pe_mode="concat",
+        num_pe_neuron=10,
+        neuron_pe_scale=10000.0,
+        dropout=0.1,
+        num_steps=4,
+    ):
+        super().__init__()
+        self.max_len = 50000  # different from windows
+        self.pe_mode = pe_mode
+        self.neuron_pe_scale = neuron_pe_scale
+        self.dropout = nn.Dropout(p=dropout)
+        if self.pe_mode == "concat":
+            self.num_pe_neuron = copy.deepcopy(num_pe_neuron)
+        elif self.pe_mode == "add":
+            self.num_pe_neuron = copy.deepcopy(d_model)
+        pe = torch.zeros(self.max_len, self.num_pe_neuron)  # MaxL, Neur
+        position = torch.arange(0, self.max_len, dtype=torch.float).unsqueeze(
+            1
+        )  # MaxL, 1
+        div_term = torch.exp(
+            torch.arange(0, self.num_pe_neuron, 2).float()
+            * (-math.log(neuron_pe_scale) / self.num_pe_neuron)
+        )
+        div_term_single = torch.exp(
+            torch.arange(0, self.num_pe_neuron - 1, 2).float()
+            * (-math.log(neuron_pe_scale) / self.num_pe_neuron)
+        )
+        pe[:, 0::2] = torch.heaviside(
+            torch.sin(position * div_term) - 0.8, torch.tensor([1.0])
+        )
+        pe[:, 1::2] = torch.heaviside(
+            torch.cos(position * div_term_single) - 0.8, torch.tensor([1.0])
+        )
+        pe = pe.unsqueeze(0).transpose(0, 1)  # MaxL, 1, Neur
+        print("pe.shape: ", pe.shape)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        # T, B, L, D
+        T, B, L, _ = x.shape
+        x = x.permute(1, 0, 2, 3)  # B, T, L, D
+        x = x.flatten(1, 2)  # B, TL, D
+        if self.pe_mode == "concat":
+            # tmp: TL, 1, Neur -> TL, B, Neur -> B, TL, Neur
+            tmp = self.pe[: x.size(-2), :].repeat(1, B, 1).transpose(0, 1)
+            x = torch.concat([x, tmp], dim=-1)
+            # print(x.shape) # B, TL, D'
+        elif self.pe_mode == "add":
+            # [B, TL, D] + [1, TL, Neur]
+            # print(self.pe[:x.size(-2), :].shape)
+            x = x + self.pe[: x.size(-2), :].transpose(0, 1)
+            # print(x.shape) # B, TL, D
+        x = x.transpose(0, 1)  # TL, B D
+        x = x.reshape(T, L, B, -1)  # T, L, B, D
+        x = x.permute(0, 2, 1, 3)  # T, B, L, D
+        return self.dropout(x)
+
+
+class StaticPE(nn.Module):
+    r"""Inject some information about the relative or absolute position of the tokens
+        in the sequence. The positional encodings have the same dimension as
+        the embeddings, so that the two can be summed. Here, we use sine and cosine
+        functions of different frequencies.
+    .. math::
+        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
+        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
+        \text{where pos is the word position and i is the embed idx)"""
+
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        pe = torch.zeros(max_len, d_model)  # MaxL, D
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # MaxL, 1
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
+        )
+        div_term_single = torch.exp(
+            torch.arange(0, d_model - 1, 2).float() * (-math.log(10000.0) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term_single)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        # x: L, TB, D
+        x = x + self.pe[: x.size(0), :]
+        x = self.dropout(x)
+        return x
+
+
+class ConvPE(nn.Module):
+    def __init__(self, d_model, dropout=0.1, max_len=5000, num_steps=4):
+
+        super().__init__()
+        self.T = num_steps
+        self.rpe_conv = nn.Conv1d(
+            d_model, d_model, kernel_size=3, stride=1, padding=1, bias=False
+        )
+        self.rpe_bn = nn.BatchNorm1d(d_model)
+        self.rpe_lif = neuron.LIFNode(
+            step_mode="m",
+            detach_reset=True,
+            surrogate_function=surrogate.ATan(),
+            v_threshold=1.0,
+        )
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x):
+        # x: L, TB, D
+        L, TB, D = x.shape
+        x_feat = x.permute(1, 2, 0)  # TB, D, L
+        x_feat = self.rpe_conv(x_feat)  # TB, D, L
+        x_feat = (
+            self.rpe_bn(x_feat).reshape(self.T, int(TB / self.T), D, L).contiguous()
+        )  # T, B, D, L
+        x_feat = self.rpe_lif(x_feat)
+        x_feat = x_feat.flatten(0, 1)  # TB, D, L
+        x_feat = self.dropout(x_feat)  # TB, D, L
+        x_feat = x_feat.permute(2, 0, 1)  # L, TB, D
+        x = x + x_feat
+        return x
+
+
+class PositionEmbedding(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        pe_type: str,
+        max_len: int = 5000,
+        pe_mode: str = "add",
+        num_pe_neuron: int = 10,
+        neuron_pe_scale: float = 1000.0,
+        dropout=0.1,
+        num_steps=4,
+    ):
+        super().__init__()
+        self.emb_type = pe_type
+        if pe_type in ["learn", "none"]:
+            self.emb = nn.Embedding(max_len, input_size)
+        elif pe_type == "conv":
+            self.emb = ConvPE(
+                d_model=input_size,
+                max_len=max_len,
+                dropout=dropout,
+                num_steps=num_steps,
+            )
+        elif pe_type == "static":
+            self.emb = StaticPE(d_model=input_size, max_len=max_len, dropout=dropout)
+        elif pe_type == "neuron":
+            self.emb = NeuronPE(
+                d_model=input_size,
+                pe_mode=pe_mode,
+                num_pe_neuron=num_pe_neuron,
+                neuron_pe_scale=neuron_pe_scale,
+                dropout=dropout,
+                num_steps=num_steps,
+            )
+        elif pe_type == "random":
+            self.emb = RandomPE(
+                d_model=input_size,
+                pe_mode=pe_mode,
+                num_pe_neuron=num_pe_neuron,
+                neuron_pe_scale=neuron_pe_scale,
+                dropout=dropout,
+                num_steps=num_steps,
+            )
+        else:
+            raise ValueError("Unknown embedding type: {}".format(pe_type))
+
+    def forward(self, x):
+        if self.emb_type == "learn":
+            # T, B, L, D = x.shape # x: T, B, L, D
+            # x = x.flatten(0, 1) # TB, L, D
+            tmp = torch.arange(
+                end=x.size()[1], device=x.device
+            )  # [0,1,2,...,L-1], shape: L
+            embedding = self.emb(tmp)  # shape: L, D
+            embedding = embedding.repeat([x.size()[0], 1, 1])  # TB, L, D'
+            x = x + embedding
+            # x = x.reshape(T, B, L, -1)
+        elif self.emb_type in ["static", "conv"]:
+            T, B, L, _ = x.shape  # x: T, B, L, D
+            x = x.flatten(0, 1)  # TB, L, D
+            x = self.emb(x.transpose(0, 1)).transpose(0, 1)  # x: TB, L, D'
+            x = x.reshape(T, B, L, -1)
+        elif self.emb_type in ["neuron", "random"]:
+            T, B, L, _ = x.shape  # x: T, B, L, D
+            # T, B, L, D
+            x = self.emb(x)
+            x = x.reshape(T, B, L, -1)
+        return x  # T, B, L, D'
+
+
+tau = 2.0  # beta = 1 - 1/tau
+backend = "torch"
+detach_reset = True
+
+
+class RepeatEncoder(nn.Module):
+    def __init__(self, output_size: int):
+        super().__init__()
+        self.out_size = output_size
+        self.lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: B, L, C
+        inputs = inputs.repeat(
+            tuple([self.out_size] + torch.ones(len(inputs.size()), dtype=int).tolist())
+        )  # T B L C
+        inputs = inputs.permute(0, 1, 3, 2)  # T B C L
+        spks = self.lif(inputs)  # T B C L
+        return spks
+
+
+class DeltaEncoder(nn.Module):
+    def __init__(self, output_size: int):
+        super().__init__()
+        self.norm = nn.BatchNorm2d(1)
+        self.enc = nn.Linear(1, output_size)
+        self.lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: B, L, C
+        delta = torch.zeros_like(inputs)
+        delta[:, 1:] = inputs[:, 1:, :] - inputs[:, :-1, :]
+        delta = delta.unsqueeze(1).permute(0, 1, 3, 2)  # B, 1, C, L
+        delta = self.norm(delta)
+        delta = delta.permute(0, 2, 3, 1)  # B, C, L, 1
+        enc = self.enc(delta)  # B, C, L, T
+        enc = enc.permute(3, 0, 1, 2)  # T, B, C, L
+        spks = self.lif(enc)
+        return spks
+
+
+class ConvEncoder(nn.Module):
+    def __init__(self, output_size: int, kernel_size: int = 3):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv2d(
+                in_channels=1,
+                out_channels=output_size,
+                kernel_size=(1, kernel_size),
+                stride=1,
+                padding=(0, kernel_size // 2),
+            ),
+            nn.BatchNorm2d(output_size),
+        )
+        self.lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: B, L, C
+        inputs = inputs.permute(0, 2, 1).unsqueeze(1)  # B, 1, C, L
+        enc = self.encoder(inputs)  # B, T, C, L
+        enc = enc.permute(1, 0, 2, 3)  # T, B, C, L
+        spks = self.lif(enc)  # T, B, C, L
+        return spks
+
+
+
+
+SpikeEncoder = {
+    "snntorch": {
+        "repeat": RepeatEncoder,
+        "conv": ConvEncoder,
+        "delta": DeltaEncoder,
+    },
+    "spikingjelly": {
+        "repeat": RepeatEncoder,
+        "conv": ConvEncoder,
+        "delta": DeltaEncoder,
+    },
+}
+
+
+
+class SSA(nn.Module):
+    def __init__(
+        self, length, tau, common_thr, dim, heads=8, qkv_bias=False, qk_scale=0.25
+    ):
+        super().__init__()
+        assert dim % heads == 0, f"dim {dim} should be divided by num_heads {heads}."
+
+        self.dim = dim
+        self.heads = heads
+        self.qk_scale = qk_scale
+
+        self.q_m = nn.Linear(dim, dim)
+        self.q_bn = nn.BatchNorm1d(dim)
+
+
+        self.q_tslif = TSLIFNode(
+            surrogate_function=SG.apply,
+        )
+
+        self.k_m = nn.Linear(dim, dim)
+        self.k_bn = nn.BatchNorm1d(dim)
+
+
+        self.k_tslif = TSLIFNode(
+            surrogate_function =SG.apply,
+        )
+
+        self.v_m = nn.Linear(dim, dim)
+        self.v_bn = nn.BatchNorm1d(dim)
+
+        self.v_tslif = TSLIFNode(
+            surrogate_function =SG.apply,
+        )
+
+
+        self.attn_tslif = TSLIFNode(
+            v_threshold=0.7,
+            surrogate_function=SG.apply
+        )
+
+        self.last_m = nn.Linear(dim, dim)
+        self.last_bn = nn.BatchNorm1d(dim)
+
+        self.last_tslif = TSLIFNode(
+            surrogate_function=SG.apply
+        )
+
+    def forward(self, x):
+        utils.reset(self.q_tslif)
+        utils.reset(self.k_tslif)
+        utils.reset(self.v_tslif)
+        utils.reset(self.attn_tslif)
+        utils.reset(self.last_tslif)
+        # x = x.transpose(0, 1)
+
+        # T, B, L, D = x.shape
+        B, T, L, D = x.shape
+        x_for_qkv = x.flatten(0, 1)  # BT L D
+        q_m_out = self.q_m(x_for_qkv) # BT L D
+
+        q_m_out = (
+            self.q_bn(q_m_out.transpose(-1, -2))
+            .transpose(-1, -2)
+            .reshape(B, T, L, D)
+            .contiguous()
+        )
+        q_m_out = self.q_tslif(q_m_out)
+
+        q = (
+            q_m_out.reshape(B, T, L, self.heads, D // self.heads)
+            .permute(0, 1, 3, 2, 4)
+            .contiguous()
+        )
+
+        k_m_out = self.k_m(x_for_qkv)
+
+        k_m_out = (
+            self.k_bn(k_m_out.transpose(-1, -2))
+            .transpose(-1, -2)
+            .reshape(B, T, L, D)
+            .contiguous()
+        )
+
+        k_m_out = self.k_tslif(k_m_out)
+        k = (
+            k_m_out.reshape(B, T, L, self.heads, D // self.heads)
+            .permute(0, 1, 3, 2, 4)
+            .contiguous()
+        )
+
+        v_m_out = self.v_m(x_for_qkv)
+        v_m_out = (
+            self.v_bn(v_m_out.transpose(-1, -2))
+            .transpose(-1, -2)
+            .reshape(B, T, L, D)
+            .contiguous()
+        )
+
+        v_m_out = self.v_tslif(v_m_out)
+
+
+        v = (
+            v_m_out.reshape(B, T, L, self.heads, D // self.heads)
+            .permute(0, 1, 3, 2, 4)
+            .contiguous()
+        )
+
+        attn = (q @ k.transpose(-2, -1)) * self.qk_scale
+        x = attn @ v  # x_shape: T * B * heads * L * D//heads
+
+        x = x.transpose(2, 3).reshape(B, T, L, D).contiguous()
+        x = self.attn_tslif(x)
+        x = x.flatten(0, 1)
+        x = self.last_m(x)
+        x = self.last_bn(x.transpose(-1, -2)).transpose(-1, -2)
+        x = self.last_tslif(x.reshape(B, T, L, D).contiguous())
+        return x
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        length,
+        tau,
+        common_thr,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        self.in_features = in_features
+        self.hidden_features = hidden_features
+        self.out_features = out_features
+
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.bn1 = nn.BatchNorm1d(hidden_features)
+
+        self.mlp_tclif1 = TCLIFNode2(
+            surrogate_function =SG.apply,
+        )
+
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.bn2 = nn.BatchNorm1d(out_features)
+
+
+
+        self.mlp_tclif2 =  TCLIFNode(
+            surrogate_function =SG.apply,
+        )
+
+    def forward(self, x):
+        utils.reset(self.mlp_tclif1)
+        utils.reset(self.mlp_tclif2)
+        # T, B, L, D = x.shape
+        B, T, L, D = x.shape
+        x = x.flatten(0, 1) # BT L D
+        x = self.fc1(x)  # TB L H
+        x = (
+            self.bn1(x.transpose(-1, -2))
+            .transpose(-1, -2)
+            .reshape(B, T, L, self.hidden_features)
+            .contiguous()
+        )
+        x = self.mlp_tclif1(x)
+        x = x.flatten(0, 1)  # TB L H
+        x = self.fc2(x)  # TB L D
+        x = (
+            self.bn2(x.transpose(-1, -2))
+            .transpose(-1, -2)
+            .reshape(B, T, L, D)
+            .contiguous()
+        )
+        x = self.mlp_tclif2(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        length,
+        tau,
+        common_thr,
+        dim,
+        d_ff,
+        heads=8,
+        qkv_bias=False,
+        qk_scale=0.125,
+    ):
+        super().__init__()
+        self.attn = SSA(
+            length=length,
+            tau=tau,
+            common_thr=common_thr,
+            dim=dim,
+            heads=heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+        )
+        self.mlp = MLP(
+            length=length,
+            tau=tau,
+            common_thr=common_thr,
+            in_features=dim,
+            hidden_features=d_ff,
+        )
+
+    def forward(self, x):
+        x = x + self.attn(x)
+        x = x + self.mlp(x)
+        return x
+
+
+
+
+
+@torch.jit.script
+def heaviside(x: torch.Tensor):
+    return (x >= 0).to(x)
+
+@torch.jit.script
+def atan_backward(grad_output: torch.Tensor, x: torch.Tensor, alpha: float):
+
+    return alpha / 2 / (1 + (math.pi / 2 * alpha * x).pow_(2)) * grad_output, None
+    #          
+
+class SG(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, alpha=2.0):
+        if x.requires_grad:
+            #ctx.save_for_backward(x.detach().clone())   # additional instead
+            ctx.save_for_backward(x)
+            ctx.alpha = alpha
+        return heaviside(x)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return atan_backward(grad_output, ctx.saved_tensors[0], ctx.alpha)
+
+
+class MemoryModule(nn.Module):
+    def __init__(self):
+        """
+        * :ref:`API in English <MemoryModule.__init__-en>`
+
+        .. _MemoryModule.__init__-cn:
+
+        ``MemoryModule`` 是SpikingJelly中所有有状态（记忆）模块的基类。
+
+        * :ref:`中文API <MemoryModule.__init__-cn>`
+
+        .. _MemoryModule.__init__-en:
+
+        ``MemoryModule`` is the base class of all stateful modules in SpikingJelly.
+
+        """
+        super().__init__()
+        self._memories = {}
+        self._memories_rv = {}
+
+    def register_memory(self, name: str, value):
+        """
+        * :ref:`API in English <MemoryModule.register_memory-en>`
+
+        .. _MemoryModule.register_memory-cn:
+
+        :param name: 变量的名字
+        :type name: str
+        :param value: 变量的值
+        :type value: any
+
+        将变量存入用于保存有状态变量（例如脉冲神经元的膜电位）的字典中。这个变量的重置值会被设置为 ``value``。每次调用 ``self.reset()``
+        函数后， ``self.name`` 都会被重置为 ``value``。
+
+        * :ref:`中文API <MemoryModule.register_memory-cn>`
+
+        .. _MemoryModule.register_memory-en:
+
+        :param name: variable's name
+        :type name: str
+        :param value: variable's value
+        :type value: any
+
+        Register the variable to memory dict, which saves stateful variables (e.g., the membrane potential of a
+        spiking neuron). The reset value of this variable will be ``value``. ``self.name`` will be set to ``value`` after
+        each calling of ``self.reset()``.
+
+        """
+        assert not hasattr(self, name), f'{name} has been set as a member variable!'
+        self._memories[name] = value
+        self.set_reset_value(name, value)
+
+    def reset(self):
+        """
+        * :ref:`API in English <MemoryModule.reset-en>`
+
+        .. _MemoryModule.reset-cn:
+
+        重置所有有状态变量为默认值。
+
+        * :ref:`中文API <MemoryModule.reset-cn>`
+
+        .. _MemoryModule.reset-en:
+
+        Reset all stateful variables to their default values.
+        """
+        for key in self._memories.keys():
+            self._memories[key] = copy.deepcopy(self._memories_rv[key])
+
+    def set_reset_value(self, name: str, value):
+        self._memories_rv[name] = copy.deepcopy(value)
+
+    def __getattr__(self, name: str):
+        if '_memories' in self.__dict__:
+            memories = self.__dict__['_memories']
+            if name in memories:
+                return memories[name]
+
+        return super().__getattr__(name)
+
+    def __setattr__(self, name: str, value) -> None:
+        _memories = self.__dict__.get('_memories')
+        if _memories is not None and name in _memories:
+            _memories[name] = value
+        else:
+            super().__setattr__(name, value)
+
+    def __delattr__(self, name):
+        if name in self._memories:
+            del self._memories[name]
+            del self._memories_rv[name]
+        else:
+            return super().__delattr__(name)
+
+    def __dir__(self):
+        module_attrs = dir(self.__class__)
+        attrs = list(self.__dict__.keys())
+        parameters = list(self._parameters.keys())
+        modules = list(self._modules.keys())
+        buffers = list(self._buffers.keys())
+        memories = list(self._memories.keys())
+        keys = module_attrs + attrs + parameters + modules + buffers + memories
+
+        # Eliminate attrs that are not legal Python variable names
+        keys = [key for key in keys if not key[0].isdigit()]
+
+        return sorted(keys)
+
+    def memories(self):
+        """
+        * :ref:`API in English <MemoryModule.memories-en>`
+
+        .. _MemoryModule.memories-cn:
+
+        :return: 返回一个所有状态变量的迭代器
+        :rtype: Iterator
+
+        * :ref:`中文API <MemoryModule.memories-cn>`
+
+        .. _MemoryModule.memories-en:
+
+        :return: an iterator over all stateful variables
+        :rtype: Iterator
+        """
+        for name, value in self._memories.items():
+            yield value
+
+    def named_memories(self):
+        """
+        * :ref:`API in English <MemoryModule.named_memories-en>`
+
+        .. _MemoryModule.named_memories-cn:
+
+        :return: 返回一个所有状态变量及其名称的迭代器
+        :rtype: Iterator
+
+        * :ref:`中文API <MemoryModule.named_memories-cn>`
+
+        .. _MemoryModule.named_memories-en:
+
+        :return: an iterator over all stateful variables and their names
+        :rtype: Iterator
+        """
+
+        for name, value in self._memories.items():
+            yield name, value
+
+    def detach(self):
+        """
+        * :ref:`API in English <MemoryModule.detach-en>`
+
+        .. _MemoryModule.detach-cn:
+
+        从计算图中分离所有有状态变量。
+
+        .. tip::
+
+            可以使用这个函数实现TBPTT(Truncated Back Propagation Through Time)。
+
+
+        * :ref:`中文API <MemoryModule.detach-cn>`
+
+        .. _MemoryModule.detach-en:
+
+        Detach all stateful variables.
+
+        .. admonition:: Tip
+            :class: tip
+
+            We can use this function to implement TBPTT(Truncated Back Propagation Through Time).
+
+        """
+
+        for key in self._memories.keys():
+            if isinstance(self._memories[key], torch.Tensor):
+                self._memories[key].detach_()
+
+    def _apply(self, fn):
+        for key, value in self._memories.items():
+            if isinstance(value, torch.Tensor):
+                self._memories[key] = fn(value)
+        # do not apply on default values
+        # for key, value in self._memories_rv.items():
+        #     if isinstance(value, torch.Tensor):
+        #         self._memories_rv[key] = fn(value)
+        return super()._apply(fn)
+
+    def _replicate_for_data_parallel(self):
+        replica = super()._replicate_for_data_parallel()
+        replica._memories = self._memories.copy()
+        return replica
+
+
+class StepModule:
+    def supported_step_mode(self):
+        """
+        * :ref:`API in English <StepModule.supported_step_mode-en>`
+        .. _StepModule.supported_step_mode-cn:
+        :return: 包含支持的后端的tuple
+        :rtype: tuple[str]
+        返回此模块支持的步进模式。
+        * :ref:`中文 API <StepModule.supported_step_mode-cn>`
+        .. _StepModule.supported_step_mode-en:
+        :return: a tuple that contains the supported backends
+        :rtype: tuple[str]
+        """
+        return ('s', 'm')
+
+    @property
+    def step_mode(self):
+        """
+        * :ref:`API in English <StepModule.step_mode-en>`
+        .. _StepModule.step_mode-cn:
+        :return: 模块当前使用的步进模式
+        :rtype: str
+        * :ref:`中文 API <StepModule.step_mode-cn>`
+        .. _StepModule.step_mode-en:
+        :return: the current step mode of this module
+        :rtype: str
+        """
+        return self._step_mode
+
+    @step_mode.setter
+    def step_mode(self, value: str):
+        """
+        * :ref:`API in English <StepModule.step_mode-setter-en>`
+        .. _StepModule.step_mode-setter-cn:
+        :param value: 步进模式
+        :type value: str
+        将本模块的步进模式设置为 ``value``
+        * :ref:`中文 API <StepModule.step_mode-setter-cn>`
+        .. _StepModule.step_mode-setter-en:
+        :param value: the step mode
+        :type value: str
+        Set the step mode of this module to be ``value``
+        """
+        if value not in self.supported_step_mode():
+            raise ValueError(f'step_mode can only be {self.supported_step_mode()}, but got "{value}"!')
+        self._step_mode = value
+
+
+
+class BaseNode(MemoryModule):
+    def __init__(self,
+                 v_threshold: float = 1.,
+                 v_reset: float = 0.,
+                 surrogate_function: Callable = None,
+                 detach_reset: bool = False,
+                 step_mode='s', backend='torch',
+                 store_v_seq: bool = True):
+
+        assert isinstance(v_reset, float) or v_reset is None
+        assert isinstance(v_threshold, float)
+        assert isinstance(detach_reset, bool)
+        super().__init__()
+
+        if v_reset is None:
+            self.register_memory('v', 0.)
+            self.register_memory('v_s', 0.)
+        else:
+            self.register_memory('v', v_reset)
+
+        self.v_threshold = v_threshold
+
+        self.v_reset = v_reset
+        self.detach_reset = detach_reset
+        self.surrogate_function = surrogate_function
+
+        self.step_mode = step_mode
+        self.backend = backend
+
+        self.store_v_seq = store_v_seq
+
+
+        self.alpha_s = torch.nn.Parameter(torch.tensor(0.5, dtype=torch.float)) 
+        self.alpha_l = torch.nn.Parameter(torch.tensor(0.5, dtype=torch.float))
+        #self.alpha_s = torch.nn.Parameter(torch.randn([1, 128], dtype=torch.float))
+        #self.alpha_l = torch.nn.Parameter(torch.randn([1, 128], dtype=torch.float))
+
+    @property
+    def store_v_seq(self):
+        return self._store_v_seq
+
+    @store_v_seq.setter
+    def store_v_seq(self, value: bool):
+        self._store_v_seq = value
+        if value:
+            if not hasattr(self, 'v_seq'):
+                self.register_memory('v_seq', None)
+
+    @staticmethod
+    @torch.jit.script
+    def jit_hard_reset(v: torch.Tensor, spike: torch.Tensor, v_reset: float):
+        v = (1. - spike) * v + spike * v_reset
+
+        return v
+
+    @staticmethod
+    @torch.jit.script
+    def jit_soft_reset(v: torch.Tensor, spike: torch.Tensor, v_threshold: float):
+        v = v - spike * v_threshold
+        return v
+
+
+    @abstractmethod
+    def neuronal_charge(self, x: torch.Tensor):
+        raise NotImplementedError
+
+    def neuronal_fire(self):
+        return self.surrogate_function(self.v - self.v_threshold, 2.0)
+
+    def sl_neuronal_fire(self):
+        s_s = self.surrogate_function(self.v - self.v_threshold, 2.0)
+        s_l = self.surrogate_function(self.v_s - self.v_threshold,  2.0)
+        return s_s, s_l
+
+    def extra_repr(self):
+        return f'v_threshold={self.v_threshold}, v_reset={self.v_reset}, detach_reset={self.detach_reset}, step_mode={self.step_mode}, backend={self.backend}'
+
+    def single_step_forward(self, x: torch.Tensor):
+        self.v_float_to_tensor(x)
+        self.neuronal_charge(x)
+        s_s, s_l = self.sl_neuronal_fire()
+        spike = self.alpha_s * s_s + self.alpha_l * s_l
+        self.neuronal_reset(s_s, s_l)
+        
+        return spike
+
+    def multi_step_forward(self, x_seq: torch.Tensor):
+
+        #### time series ###
+        T = x_seq.shape[-1]
+        y_seq = []
+        if self.store_v_seq:
+            v_seq = []
+        for t in range(T):
+            y = self.single_step_forward(x_seq[:, t])
+            y_seq.append(y)
+            if self.store_v_seq:
+                v_seq.append(self.v)
+        if self.store_v_seq:
+            self.v_seq = torch.stack(v_seq)
+
+        # if self.store_v_seq:
+        #     self.v_seq = torch.stack(v_seq)
+        outputs = torch.stack(y_seq, dim=0).permute(1, 0)
+
+        return outputs
+
+    def v_float_to_tensor(self, x: torch.Tensor):
+        if isinstance(self.v, float):
+            v_init = self.v
+            self.v = torch.full_like(x.data, v_init)
+
+
+class TSLIFNode(BaseNode):
+    def __init__(self,
+                 v_threshold=1.0,
+                 v_reset=0.,
+                 surrogate_function: Callable = None,
+                 detach_reset=False,
+                 hard_reset=False,
+                 step_mode='s',
+                 k=2,
+                 decay_factor: torch.Tensor = torch.tensor([0.8, 0.2, 0.3, 0.7], dtype=torch.float),
+                 gamma: float = 0.5):
+        super(TSLIFNode, self).__init__(v_threshold, v_reset, surrogate_function, detach_reset, step_mode)
+        self.k = k
+        for i in range(1, self.k + 1):
+            self.register_memory('v' + str(i), 0.)
+
+
+        self.names = self._memories
+        self.hard_reset = hard_reset
+        self.gamma = gamma
+        self.decay_factor = torch.nn.Parameter(decay_factor)
+        self.kk = torch.nn.Parameter(torch.tensor([0.8], dtype=torch.float))
+        self.yy = torch.nn.Parameter(torch.tensor([0.1], dtype=torch.float))
+
+    @property
+    def supported_backends(self):
+        if self.step_mode == 's':
+            return ('torch',)
+        elif self.step_mode == 'm':
+            return ('torch', 'cupy')
+        else:
+            raise ValueError(self.step_mode)
+
+    def neuronal_charge(self, x: torch.Tensor):
+        self.names['v1'] = self.decay_factor[0] * self.names['v1'] + self.decay_factor[1] * x - self.yy * self.names['v2']
+        self.names['v2'] = self.decay_factor[2] * self.names['v2'] + self.decay_factor[3] * x - self.kk * self.names['v1']
+        self.v = self.names['v2']
+        self.v_s = self.names['v1']
+
+    def neuronal_reset(self, spike_s, spike_l):
+        if not self.hard_reset:
+            self.names['v1'] = self.jit_soft_reset(self.names['v1'], spike_l, self.gamma)
+            self.names['v2'] = self.jit_soft_reset(self.names['v2'], spike_s, self.v_threshold)
+        else:
+            for i in range(2, self.k + 1):
+                self.names['v' + str(i)] = self.jit_hard_reset(self.names['v' + str(i)], spike_s, self.v_reset)
+
+    def forward(self, x: torch.Tensor):
+        return super().single_step_forward(x)
+    def extra_repr(self):
+        return f"v_threshold={self.v_threshold}, v_reset={self.v_reset}, detach_reset={self.detach_reset}, " \
+               f"hard_reset={self.hard_reset}, " \
+               f"gamma={self.gamma}, k={self.k}, step_mode={self.step_mode}, backend={self.backend}"
+
+
+
+
+
+class BaseNode1(MemoryModule):
+    def __init__(self,
+                 v_threshold: float = 1.,
+                 v_reset: float = 0.,
+                 surrogate_function: Callable = None,
+                 detach_reset: bool = False,
+                 step_mode='s', backend='torch',
+                 store_v_seq: bool = True):
+
+        assert isinstance(v_reset, float) or v_reset is None
+        assert isinstance(v_threshold, float)
+        assert isinstance(detach_reset, bool)
+        super().__init__()
+
+        if v_reset is None:
+            self.register_memory('v', 0.)
+            self.register_memory('v_s', 0.)
+        else:
+            self.register_memory('v', v_reset)
+
+        self.v_threshold = v_threshold
+
+        self.v_reset = v_reset
+        self.detach_reset = detach_reset
+        self.surrogate_function = surrogate_function
+
+        self.step_mode = step_mode
+        self.backend = backend
+
+        self.store_v_seq = store_v_seq
+        self.alpha_s = torch.nn.Parameter(torch.tensor(0.5, dtype=torch.float)) 
+        self.alpha_l = torch.nn.Parameter(torch.tensor(0.5, dtype=torch.float))
+
+    @property
+    def store_v_seq(self):
+        return self._store_v_seq
+
+    @store_v_seq.setter
+    def store_v_seq(self, value: bool):
+        self._store_v_seq = value
+        if value:
+            if not hasattr(self, 'v_seq'):
+                self.register_memory('v_seq', None)
+
+    @staticmethod
+    @torch.jit.script
+    def jit_hard_reset(v: torch.Tensor, spike: torch.Tensor, v_reset: float):
+        v = (1. - spike) * v + spike * v_reset
+
+        return v
+
+    @staticmethod
+    @torch.jit.script
+    def jit_soft_reset(v: torch.Tensor, spike: torch.Tensor, v_threshold: float):
+        v = v - spike * v_threshold
+        return v
+
+
+    @abstractmethod
+    def neuronal_charge(self, x: torch.Tensor):
+        raise NotImplementedError
+
+    def neuronal_fire(self):
+        return self.surrogate_function(self.v - self.v_threshold, 2.0)
+
+    def sl_neuronal_fire(self):
+        s_s = self.surrogate_function(self.v - self.v_threshold, 2.0)
+        s_l = self.surrogate_function(self.v_s - self.v_threshold,  2.0)
+        return s_s, s_l
+
+    def extra_repr(self):
+        return f'v_threshold={self.v_threshold}, v_reset={self.v_reset}, detach_reset={self.detach_reset}, step_mode={self.step_mode}, backend={self.backend}'
+
+    def single_step_forward(self, x: torch.Tensor):
+        self.v_float_to_tensor(x)
+        self.neuronal_charge(x)
+        s_s, s_l = self.sl_neuronal_fire()
+        spike = self.alpha_s * s_s + self.alpha_l * s_l
+        self.neuronal_reset(s_s, s_l)
+        return spike
+
+    def multi_step_forward(self, x_seq: torch.Tensor):
+
+        #### time series ###
+        T = x_seq.shape[-1]
+        y_seq = []
+        if self.store_v_seq:
+            v_seq = []
+        for t in range(2):
+            y = self.single_step_forward(x_seq[:, t, :, :])
+            y_seq.append(y)
+            if self.store_v_seq:
+                v_seq.append(self.v)
+        if self.store_v_seq:
+            self.v_seq = torch.stack(v_seq)
+        outputs = torch.stack(y_seq, dim=0)
+        outputs = outputs.permute(1, 0, 2, 3)
+
+        return outputs
+
+
+    def v_float_to_tensor(self, x: torch.Tensor):
+        if isinstance(self.v, float):
+            v_init = self.v
+            self.v = torch.full_like(x.data, v_init)
+
+
+
+class TCLIFNode2(BaseNode1):
+    def __init__(self,
+                 v_threshold=0.8,
+                 v_reset=0.,
+                 surrogate_function: Callable = None,
+                 detach_reset=False,
+                 hard_reset=False,
+                 step_mode='s',
+                 k=2,
+                 decay_factor: torch.Tensor = torch.tensor([0.8, 0.2, 0.3, 0.7], dtype=torch.float),
+                 gamma: float = 0.5):
+        super(TCLIFNode2, self).__init__(v_threshold, v_reset, surrogate_function, detach_reset, step_mode)
+        self.k = k
+        for i in range(1, self.k + 1):
+            self.register_memory('v' + str(i), 0.)
+
+        self.names = self._memories
+        self.hard_reset = hard_reset
+        self.gamma = gamma
+        self.decay_factor = torch.nn.Parameter(decay_factor)
+        self.kk = torch.nn.Parameter(torch.tensor([0.8], dtype=torch.float))
+        self.yy = torch.nn.Parameter(torch.tensor([0.1], dtype=torch.float))
+
+    @property
+    def supported_backends(self):
+        if self.step_mode == 's':
+            return ('torch',)
+        elif self.step_mode == 'm':
+            return ('torch', 'cupy')
+        else:
+            raise ValueError(self.step_mode)
+
+    def neuronal_charge(self, x: torch.Tensor):
+        self.names['v1'] = self.decay_factor[0] * self.names['v1'] + self.decay_factor[1] * x - self.yy * self.names['v2']
+        self.names['v2'] = self.decay_factor[2] * self.names['v2'] + self.decay_factor[3] * x - self.kk * self.names['v1']
+        self.v = self.names['v2']
+        self.v_s = self.names['v1']
+
+    def neuronal_reset(self, spike_s, spike_l):
+        if not self.hard_reset:
+            self.names['v1'] = self.jit_soft_reset(self.names['v1'], spike_l , self.gamma)
+            self.names['v2'] = self.jit_soft_reset(self.names['v2'], spike_s, self.v_threshold)
+        else:
+            # hard reset
+            for i in range(2, self.k + 1):
+                self.names['v' + str(i)] = self.jit_hard_reset(self.names['v' + str(i)], spike_d,  self.v_reset)
+
+    def forward(self, x: torch.Tensor):
+        return super().single_step_forward(x)
+
+    def extra_repr(self):
+         return f"v_threshold={self.v_threshold}, v_reset={self.v_reset}, detach_reset={self.detach_reset}, " \
+                f"hard_reset={self.hard_reset}, " \
+                f"gamma={self.gamma}, k={self.k}, step_mode={self.step_mode}, backend={self.backend}"
+
+
+
+
+
+class TCLIFNode(BaseNode):
+    def __init__(self,
+                 v_threshold=1.0,
+                 v_reset=0.,
+                 surrogate_function: Callable = None,
+                 detach_reset=False,
+                 hard_reset=False,
+                 step_mode='s',
+                 k=2,
+                 decay_factor: torch.Tensor = torch.tensor([0.8, 0.2, 0.3, 0.7], dtype=torch.float),
+                 gamma: float = 0.5):
+        super(TCLIFNode, self).__init__(v_threshold, v_reset, surrogate_function, detach_reset, step_mode)
+        self.k = k
+        for i in range(1, self.k + 1):
+            self.register_memory('v' + str(i), 0.)
+
+        self.names = self._memories
+        self.hard_reset = hard_reset
+        self.gamma = gamma
+        self.decay_factor = torch.nn.Parameter(decay_factor)
+        self.kk = torch.nn.Parameter(torch.tensor([0.8], dtype=torch.float))
+        self.yy = torch.nn.Parameter(torch.tensor([0.1], dtype=torch.float))
+
+    @property
+    def supported_backends(self):
+        if self.step_mode == 's':
+            return ('torch',)
+        elif self.step_mode == 'm':
+            return ('torch', 'cupy')
+        else:
+            raise ValueError(self.step_mode)
+
+    def neuronal_charge(self, x: torch.Tensor):
+        self.names['v1'] = self.decay_factor[0] * self.names['v1'] + self.decay_factor[1] * x - self.yy * self.names['v2']
+        self.names['v2'] = self.decay_factor[2] * self.names['v2'] + self.decay_factor[3] * x - self.kk * self.names['v1']
+        self.v = self.names['v2']
+        self.v_s = self.names['v1']
+
+    def neuronal_reset(self, spike_s, spike_l):
+        if not self.hard_reset:
+            self.names['v1'] = self.jit_soft_reset(self.names['v1'], spike_l , self.gamma)
+            self.names['v2'] = self.jit_soft_reset(self.names['v2'], spike_s, self.v_threshold)
+        else:
+            # hard reset
+            for i in range(2, self.k + 1):
+                self.names['v' + str(i)] = self.jit_hard_reset(self.names['v' + str(i)], spike_d,  self.v_reset)
+
+    def forward(self, x: torch.Tensor):
+        return super().single_step_forward(x)
+    def extra_repr(self):
+        return f"v_threshold={self.v_threshold}, v_reset={self.v_reset}, detach_reset={self.detach_reset}, " \
+               f"hard_reset={self.hard_reset}, " \
+               f"gamma={self.gamma}, k={self.k}, step_mode={self.step_mode}, backend={self.backend}"
+
+
+
+
+
+class TSFormer(nn.Module):
+
+    def __init__(
+        self,
+        args,
+        dim: int = 256,
+        d_ff: Optional[int] = None,
+        num_pe_neuron: int = 40,
+        pe_type: str = "neuron",
+        pe_mode: str = "concat",  # "add" or concat
+        neuron_pe_scale: float = 10000.0,  # "100" or "1000" or "10000"
+        depths: int = 2,
+        common_thr: float = 1.0,
+        max_length: int = 5000,
+        num_steps: int = 4,
+        heads: int = 8,
+        qkv_bias: bool = False,
+        qk_scale: float = 0.125,
+        input_size: Optional[int] = None,
+        weight_file: Optional[Path] = None,
+    ):
+        super().__init__()
+        self.dim = 256
+        self.d_ff = 1024
+        self.T = args.T
+        self.depths = args.blocks
+        self.pe_type = pe_type
+        self.pe_mode = pe_mode
+        self.num_pe_neuron = num_pe_neuron
+        self.input_size = args.feature_size
+        self._snn_backend = "spikingjelly"
+        self.temporal_encoder = SpikeEncoder[self._snn_backend]["conv"](num_steps)
+        self.pre_length = args.pre_length
+        self.feature_size = args.feature_size
+        self.args = args
+        self.pe = PositionEmbedding(
+            pe_type=pe_type,
+            pe_mode=pe_mode,
+            neuron_pe_scale=neuron_pe_scale,
+            input_size=self.input_size,
+            max_len=max_length,
+            num_pe_neuron=self.num_pe_neuron,
+            dropout=0.1,
+            num_steps=num_steps,
+        )
+        if (self.pe_type == "neuron" and self.pe_mode == "concat") or (
+            self.pe_type == "random" and self.pe_mode == "concat"
+        ):
+            self.encoder = nn.Linear(self.input_size + num_pe_neuron, dim)
+        else:
+            self.encoder = nn.Linear(self.input_size, dim)
+
+        self.init_lif = neuron.LIFNode(
+            tau=tau,
+            step_mode="m",
+            detach_reset=detach_reset,
+            surrogate_function=surrogate.ATan(),
+            v_threshold=common_thr,
+            backend=backend,
+        )
+
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    length=max_length,
+                    tau=tau,
+                    common_thr=common_thr,
+                    dim=dim,
+                    d_ff=self.d_ff,
+                    heads=heads,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                )
+                for _ in range(depths)
+            ]
+        )
+
+        self.apply(self._init_weights)
+
+        self.fc = nn.Linear(args.seq_length*dim, args.pre_length*args.feature_size)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0.0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.weight, 1.0)
+            nn.init.constant_(m.bias, 0.0)
+
+    def forward(self, x):
+        functional.reset_net(self)
+        
+        if self.args.normalize:
+
+            mean = x.mean(dim=1, keepdim=True).detach() # shape [B, 1, D]
+            x = x - mean
+
+            std = torch.sqrt(torch.var(x, dim=1, keepdim=True, unbiased=False) + 1e-5).detach()
+            x = x / std
+
+        x = self.temporal_encoder(x)  # B L C -> T B C L
+        x = x.transpose(-2, -1)  # T B L C
+        if self.pe_type != "none":
+            x = self.pe(x)  # T B L C'
+        T, B, L, _ = x.shape
+        x = self.encoder(x.flatten(0, 1)).reshape(T, B, L, -1)  # T B L D
+        x = self.init_lif(x)
+
+        for blk in self.blocks:
+            x = blk(x)  # T B L D
+        out = x.mean(0) # B L D
+        out = self.fc(out.flatten(-2, -1)).reshape(-1, self.pre_length, self.feature_size)  # B D L -> B L D
+        if self.args.normalize:
+            out = out * std + mean  # denormalization
+        aux = {'gate_l0': torch.tensor(0.0, device=out.device)} # placeholder
+        return out, aux  # B D L -> B L D
+

model/TS_GRU.py

@@ -0,0 +1,640 @@
+from typing import Optional, Callable
+from pathlib import Path
+from spikingjelly.activation_based import surrogate as sj_surrogate
+from snntorch import utils
+import snntorch as snn
+from snntorch import surrogate
+import torch
+from torch import nn
+import numpy as np
+import copy
+import torch.nn.functional as F
+import math
+from abc import abstractmethod
+
+
+@torch.jit.script
+def heaviside(x: torch.Tensor):
+    return (x >= 0).to(x)
+
+@torch.jit.script
+def atan_backward(grad_output: torch.Tensor, x: torch.Tensor, alpha: float):
+
+    return alpha / 2 / (1 + (math.pi / 2 * alpha * x).pow_(2)) * grad_output, None
+           
+
+class SG(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, alpha=2.0):
+        if x.requires_grad:
+            ctx.save_for_backward(x)
+            ctx.alpha = alpha
+        return heaviside(x)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return atan_backward(grad_output, ctx.saved_tensors[0], ctx.alpha)
+
+
+class MemoryModule(nn.Module):
+    def __init__(self):
+        """
+        * :ref:`API in English <MemoryModule.__init__-en>`
+
+        .. _MemoryModule.__init__-cn:
+
+        ``MemoryModule`` 是SpikingJelly中所有有状态（记忆）模块的基类。
+
+        * :ref:`中文API <MemoryModule.__init__-cn>`
+
+        .. _MemoryModule.__init__-en:
+
+        ``MemoryModule`` is the base class of all stateful modules in SpikingJelly.
+
+        """
+        super().__init__()
+        self._memories = {}
+        self._memories_rv = {}
+
+    def register_memory(self, name: str, value):
+        """
+        * :ref:`API in English <MemoryModule.register_memory-en>`
+
+        .. _MemoryModule.register_memory-cn:
+
+        :param name: 变量的名字
+        :type name: str
+        :param value: 变量的值
+        :type value: any
+
+        将变量存入用于保存有状态变量（例如脉冲神经元的膜电位）的字典中。这个变量的重置值会被设置为 ``value``。每次调用 ``self.reset()``
+        函数后， ``self.name`` 都会被重置为 ``value``。
+
+        * :ref:`中文API <MemoryModule.register_memory-cn>`
+
+        .. _MemoryModule.register_memory-en:
+
+        :param name: variable's name
+        :type name: str
+        :param value: variable's value
+        :type value: any
+
+        Register the variable to memory dict, which saves stateful variables (e.g., the membrane potential of a
+        spiking neuron). The reset value of this variable will be ``value``. ``self.name`` will be set to ``value`` after
+        each calling of ``self.reset()``.
+
+        """
+        assert not hasattr(self, name), f'{name} has been set as a member variable!'
+        self._memories[name] = value
+        self.set_reset_value(name, value)
+
+    def reset(self):
+        """
+        * :ref:`API in English <MemoryModule.reset-en>`
+
+        .. _MemoryModule.reset-cn:
+
+        重置所有有状态变量为默认值。
+
+        * :ref:`中文API <MemoryModule.reset-cn>`
+
+        .. _MemoryModule.reset-en:
+
+        Reset all stateful variables to their default values.
+        """
+        for key in self._memories.keys():
+            self._memories[key] = copy.deepcopy(self._memories_rv[key])
+
+    def set_reset_value(self, name: str, value):
+        self._memories_rv[name] = copy.deepcopy(value)
+
+    def __getattr__(self, name: str):
+        if '_memories' in self.__dict__:
+            memories = self.__dict__['_memories']
+            if name in memories:
+                return memories[name]
+
+        return super().__getattr__(name)
+
+    def __setattr__(self, name: str, value) -> None:
+        _memories = self.__dict__.get('_memories')
+        if _memories is not None and name in _memories:
+            _memories[name] = value
+        else:
+            super().__setattr__(name, value)
+
+    def __delattr__(self, name):
+        if name in self._memories:
+            del self._memories[name]
+            del self._memories_rv[name]
+        else:
+            return super().__delattr__(name)
+
+    def __dir__(self):
+        module_attrs = dir(self.__class__)
+        attrs = list(self.__dict__.keys())
+        parameters = list(self._parameters.keys())
+        modules = list(self._modules.keys())
+        buffers = list(self._buffers.keys())
+        memories = list(self._memories.keys())
+        keys = module_attrs + attrs + parameters + modules + buffers + memories
+        keys = [key for key in keys if not key[0].isdigit()]
+
+        return sorted(keys)
+
+    def memories(self):
+        """
+        * :ref:`API in English <MemoryModule.memories-en>`
+
+        .. _MemoryModule.memories-cn:
+
+        :return: 返回一个所有状态变量的迭代器
+        :rtype: Iterator
+
+        * :ref:`中文API <MemoryModule.memories-cn>`
+
+        .. _MemoryModule.memories-en:
+
+        :return: an iterator over all stateful variables
+        :rtype: Iterator
+        """
+        for name, value in self._memories.items():
+            yield value
+
+    def named_memories(self):
+        """
+        * :ref:`API in English <MemoryModule.named_memories-en>`
+
+        .. _MemoryModule.named_memories-cn:
+
+        :return: 返回一个所有状态变量及其名称的迭代器
+        :rtype: Iterator
+
+        * :ref:`中文API <MemoryModule.named_memories-cn>`
+
+        .. _MemoryModule.named_memories-en:
+
+        :return: an iterator over all stateful variables and their names
+        :rtype: Iterator
+        """
+
+        for name, value in self._memories.items():
+            yield name, value
+
+    def detach(self):
+        """
+        * :ref:`API in English <MemoryModule.detach-en>`
+
+        .. _MemoryModule.detach-cn:
+
+        从计算图中分离所有有状态变量。
+
+        .. tip::
+
+            可以使用这个函数实现TBPTT(Truncated Back Propagation Through Time)。
+
+
+        * :ref:`中文API <MemoryModule.detach-cn>`
+
+        .. _MemoryModule.detach-en:
+
+        Detach all stateful variables.
+
+        .. admonition:: Tip
+            :class: tip
+
+            We can use this function to implement TBPTT(Truncated Back Propagation Through Time).
+
+        """
+
+        for key in self._memories.keys():
+            if isinstance(self._memories[key], torch.Tensor):
+                self._memories[key].detach_()
+
+    def _apply(self, fn):
+        for key, value in self._memories.items():
+            if isinstance(value, torch.Tensor):
+                self._memories[key] = fn(value)
+        return super()._apply(fn)
+
+    def _replicate_for_data_parallel(self):
+        replica = super()._replicate_for_data_parallel()
+        replica._memories = self._memories.copy()
+        return replica
+
+
+class StepModule:
+    def supported_step_mode(self):
+        """
+        * :ref:`API in English <StepModule.supported_step_mode-en>`
+        .. _StepModule.supported_step_mode-cn:
+        :return: 包含支持的后端的tuple
+        :rtype: tuple[str]
+        返回此模块支持的步进模式。
+        * :ref:`中文 API <StepModule.supported_step_mode-cn>`
+        .. _StepModule.supported_step_mode-en:
+        :return: a tuple that contains the supported backends
+        :rtype: tuple[str]
+        """
+        return ('s', 'm')
+
+    @property
+    def step_mode(self):
+        """
+        * :ref:`API in English <StepModule.step_mode-en>`
+        .. _StepModule.step_mode-cn:
+        :return: 模块当前使用的步进模式
+        :rtype: str
+        * :ref:`中文 API <StepModule.step_mode-cn>`
+        .. _StepModule.step_mode-en:
+        :return: the current step mode of this module
+        :rtype: str
+        """
+        return self._step_mode
+
+    @step_mode.setter
+    def step_mode(self, value: str):
+        """
+        * :ref:`API in English <StepModule.step_mode-setter-en>`
+        .. _StepModule.step_mode-setter-cn:
+        :param value: 步进模式
+        :type value: str
+        将本模块的步进模式设置为 ``value``
+        * :ref:`中文 API <StepModule.step_mode-setter-cn>`
+        .. _StepModule.step_mode-setter-en:
+        :param value: the step mode
+        :type value: str
+        Set the step mode of this module to be ``value``
+        """
+        if value not in self.supported_step_mode():
+            raise ValueError(f'step_mode can only be {self.supported_step_mode()}, but got "{value}"!')
+        self._step_mode = value
+
+
+
+class BaseNode(MemoryModule):
+    def __init__(self,
+                 v_threshold: float = 1.,
+                 v_reset: float = 0.,
+                 surrogate_function: Callable = None,
+                 detach_reset: bool = False,
+                 step_mode='s', backend='torch',
+                 store_v_seq: bool = True):
+
+        assert isinstance(v_reset, float) or v_reset is None
+        assert isinstance(v_threshold, float)
+        assert isinstance(detach_reset, bool)
+        super().__init__()
+
+        if v_reset is None:
+            self.register_memory('v', 0.)
+            self.register_memory('v_s', 0.)
+        else:
+            self.register_memory('v', v_reset)
+
+        self.v_threshold = v_threshold
+
+        self.v_reset = v_reset
+        self.detach_reset = detach_reset
+        self.surrogate_function = surrogate_function
+
+        self.step_mode = step_mode
+        self.backend = backend
+
+        self.store_v_seq = store_v_seq
+        self.alpha_s = torch.nn.Parameter(torch.tensor(0.5, dtype=torch.float)) 
+        self.alpha_l = torch.nn.Parameter(torch.tensor(0.5, dtype=torch.float))
+
+    @property
+    def store_v_seq(self):
+        return self._store_v_seq
+
+    @store_v_seq.setter
+    def store_v_seq(self, value: bool):
+        self._store_v_seq = value
+        if value:
+            if not hasattr(self, 'v_seq'):
+                self.register_memory('v_seq', None)
+
+    @staticmethod
+    @torch.jit.script
+    def jit_hard_reset(v: torch.Tensor, spike: torch.Tensor, v_reset: float):
+        v = (1. - spike) * v + spike * v_reset
+
+        return v
+
+    @staticmethod
+    @torch.jit.script
+    def jit_soft_reset(v: torch.Tensor, spike: torch.Tensor, v_threshold: float):
+        v = v - spike * v_threshold
+        return v
+
+
+    @abstractmethod
+    def neuronal_charge(self, x: torch.Tensor):
+        raise NotImplementedError
+
+    def neuronal_fire(self):
+        return self.surrogate_function(self.v - self.v_threshold, 2.0)
+
+    def sl_neuronal_fire(self):
+        s_s = self.surrogate_function(self.v - self.v_threshold, 2.0)
+        s_l = self.surrogate_function(self.v_s - self.v_threshold,  2.0)
+        return s_s, s_l
+
+    def extra_repr(self):
+        return f'v_threshold={self.v_threshold}, v_reset={self.v_reset}, detach_reset={self.detach_reset}, step_mode={self.step_mode}, backend={self.backend}'
+
+    def single_step_forward(self, x: torch.Tensor):
+        self.v_float_to_tensor(x)
+        self.neuronal_charge(x)
+        s_s, s_l = self.sl_neuronal_fire()
+        spike = self.alpha_s * s_s + self.alpha_l * s_l
+        self.neuronal_reset(s_s, s_l)
+        
+        return spike
+
+    def multi_step_forward(self, x_seq: torch.Tensor):
+
+        T = x_seq.shape[-1]
+        y_seq = []
+        if self.store_v_seq:
+            v_seq = []
+        for t in range(T):
+            y = self.single_step_forward(x_seq[:, t])
+            y_seq.append(y)
+            if self.store_v_seq:
+                v_seq.append(self.v)
+        if self.store_v_seq:
+            self.v_seq = torch.stack(v_seq)
+        outputs = torch.stack(y_seq, dim=0).permute(1, 0)
+
+        return outputs
+
+    def v_float_to_tensor(self, x: torch.Tensor):
+        if isinstance(self.v, float):
+            v_init = self.v
+            self.v = torch.full_like(x.data, v_init)
+
+
+class TSLIFNode(BaseNode):
+    def __init__(self,
+                 v_threshold=1.0,
+                 v_reset=0.,
+                 surrogate_function: Callable = None,
+                 detach_reset=False,
+                 hard_reset=False,
+                 step_mode='s',
+                 k=2,
+                 decay_factor: torch.Tensor = torch.tensor([0.8, 0.2, 0.3, 0.7], dtype=torch.float),
+                 gamma: float = 0.5):
+        super(TSLIFNode, self).__init__(v_threshold, v_reset, surrogate_function, detach_reset, step_mode)
+        self.k = k
+        for i in range(1, self.k + 1):
+            self.register_memory('v' + str(i), 0.)
+        self.names = self._memories
+        self.hard_reset = hard_reset
+        self.gamma = gamma
+        self.decay_factor = torch.nn.Parameter(decay_factor)
+        self.kk = torch.nn.Parameter(torch.tensor([0.8], dtype=torch.float))
+        self.yy = torch.nn.Parameter(torch.tensor([0.1], dtype=torch.float))
+
+    @property
+    def supported_backends(self):
+        if self.step_mode == 's':
+            return ('torch',)
+        elif self.step_mode == 'm':
+            return ('torch', 'cupy')
+        else:
+            raise ValueError(self.step_mode)
+
+    def neuronal_charge(self, x: torch.Tensor):
+        self.names['v1'] = self.decay_factor[0] * self.names['v1'] + self.decay_factor[1] * x - self.yy * self.names['v2']
+        self.names['v2'] = self.decay_factor[2] * self.names['v2'] + self.decay_factor[3] * x - self.kk * self.names['v1']
+        self.v = self.names['v2']
+        self.v_s = self.names['v1']
+
+    def neuronal_reset(self, spike_s, spike_l):
+        if not self.hard_reset:
+            self.names['v1'] = self.jit_soft_reset(self.names['v1'], spike_l, self.gamma)
+            self.names['v2'] = self.jit_soft_reset(self.names['v2'], spike_s, self.v_threshold)
+        else:
+            for i in range(2, self.k + 1):
+                self.names['v' + str(i)] = self.jit_hard_reset(self.names['v' + str(i)], spike_s, self.v_reset)
+
+    def forward(self, x: torch.Tensor):
+        return super().single_step_forward(x)
+    def extra_repr(self):
+        return f"v_threshold={self.v_threshold}, v_reset={self.v_reset}, detach_reset={self.detach_reset}, " \
+               f"hard_reset={self.hard_reset}, " \
+               f"gamma={self.gamma}, k={self.k}, step_mode={self.step_mode}, backend={self.backend}"
+
+
+
+
+class GRUCell(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        num_steps: int = 4,
+        grad_slope: float = 25.0,
+        beta: float = 0.99,
+        output_mems: bool = False,
+    ):
+        super().__init__()
+        self.spike_grad = surrogate.atan(alpha=2.0)
+        self.input_size = input_size
+        self.num_steps = num_steps
+        self.hidden_size = hidden_size
+        self.beta = beta
+        self.full_rec = output_mems
+
+
+        self.linear_ih = nn.Linear(input_size, 3 * hidden_size)
+        self.linear_hh = nn.Linear(hidden_size, 3 * hidden_size)
+        self.surrogate_function1 = sj_surrogate.ATan()
+
+        self.tslif = TSLIFNode(
+            surrogate_function=SG.apply    
+        )
+
+    def forward(self, inputs):
+        if inputs.size(-1) == self.input_size:
+            h = torch.zeros(
+                size=[inputs.shape[0], self.hidden_size],
+                dtype=torch.float,
+                device=inputs.device,
+            )
+            y_ih = torch.split(self.linear_ih(inputs), self.hidden_size, dim=1)
+            y_hh = torch.split(self.linear_hh(h), self.hidden_size, dim=1)
+            r = self.surrogate_function1(y_ih[0] + y_hh[0])
+            z = self.surrogate_function1(y_ih[1] + y_hh[1])
+            n = self.surrogate_function1(y_ih[2] + r * y_hh[2])
+            h = (1.0 - z) * n + z * h
+            cur = h
+        elif inputs.size(-1) == self.num_steps and inputs.size(-2) == self.input_size:
+            inputs = inputs.transpose(-1, -2)  # BC, T, H
+            h = torch.zeros(
+                size=[inputs.shape[0], self.hidden_size, self.num_steps],
+                dtype=torch.float,
+                device=inputs.device,
+            )
+            y_ih = torch.split(
+                self.linear_ih(inputs).transpose(-1, -2), self.hidden_size, dim=1
+            )
+            y_hh = torch.split(
+                self.linear_hh(h.transpose(-1, -2)).transpose(-1, -2),
+                self.hidden_size,
+                dim=1,
+            )
+            r = self.surrogate_function1(y_ih[0] + y_hh[0])
+            z = self.surrogate_function1(y_ih[1] + y_hh[1])
+            n = self.surrogate_function1(y_ih[2] + r * y_hh[2])
+            h = (1.0 - z) * n + z * h
+            cur = h
+            static = False
+        else:
+            raise ValueError(
+                f"Input size mismatch! Got {inputs.size()} but expected "
+                f"(..., {self.input_size}, {self.num_steps}) or (..., {self.input_size})"
+            )
+
+        spks = self.tslif(cur)
+        return spks
+
+
+class DeltaEncoder(nn.Module):
+    def __init__(self, output_size: int):
+        super().__init__()
+        self.norm = nn.BatchNorm2d(1)
+        self.enc = nn.Linear(1, output_size)
+        self.lif = snn.Leaky(
+            beta=0.99, spike_grad=SG.apply, init_hidden=True, output=False
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: batch, L, C
+        delta = torch.zeros_like(inputs)
+        delta[:, 1:] = inputs[:, 1:, :] - inputs[:, :-1, :]
+        delta = delta.unsqueeze(1).permute(0, 1, 3, 2)  # batch, 1, C, L
+        delta = self.norm(delta)
+        delta = delta.permute(0, 2, 3, 1)  # batch, C, L, 1
+        enc = self.enc(delta)  # batch, C, L, output_size
+        enc = enc.permute(0, 3, 1, 2)  # batch, output_size, C, L
+        spks = self.lif(enc)
+        return spks
+
+
+class ConvEncoder(nn.Module):
+    def __init__(self, output_size: int, kernel_size: int = 3):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv2d(
+                in_channels=1,
+                out_channels=output_size,
+                kernel_size=(1, kernel_size),
+                stride=1,
+                padding=(0, kernel_size // 2),
+            ),
+            nn.BatchNorm2d(output_size),
+        )
+        self.lif = snn.Leaky(
+            beta=0.99,
+            spike_grad=surrogate.atan(alpha=2.0),
+            init_hidden=True,
+            output=False,
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: batch, L, C
+        inputs = inputs.permute(0, 2, 1).unsqueeze(1)  # batch, 1, C, L
+        enc = self.encoder(inputs)  # batch, output_size, C, L
+        spks = self.lif(enc)
+        return spks
+
+
+class TSGRU(nn.Module):
+    def __init__(
+        self,
+        args,
+        hidden_size: int,
+        layers: int = 1,
+        num_steps: int = 50,
+        grad_slope: float = 25.0,
+        input_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        weight_file: Optional[Path] = None,
+        encoder_type: Optional[str] = "conv",
+    ):
+        super().__init__()
+
+        self.hidden_size   = args.hidden_size
+        self.num_steps   = args.T
+        self.input_size = args.feature_size
+        self.pre_length   = args.pre_length
+        self.layers       = args.blocks
+        self.args = args
+
+        if encoder_type == "conv":
+            self.encoder = ConvEncoder(self.hidden_size)
+        elif encoder_type == "delta":
+            self.encoder = DeltaEncoder(self.hidden_size)
+        else:
+            raise ValueError(f"Unknown encoder type {encoder_type}")
+
+        self.net = nn.Sequential(
+            *[
+                GRUCell(
+                    self.hidden_size,          
+                    self.hidden_size,          
+                    num_steps=self.num_steps,
+                    grad_slope=grad_slope,
+                    output_mems=(i == self.layers - 1),
+                )
+                for i in range(self.layers)
+            ]
+        )
+
+        self.__output_size = self.hidden_size
+        self.fc = nn.Linear(self.__output_size, self.pre_length)
+
+        self.to('cuda:0') 
+
+    def forward(self, inputs: torch.Tensor):
+
+        utils.reset(self.encoder)
+        for layer in self.net:
+            utils.reset(layer)
+
+
+        bs, length, c_num = inputs.size() 
+
+        if self.args.normalize:
+
+            mean = inputs.mean(dim=1, keepdim=True).detach() # shape [B, 1, D]
+            inputs = inputs - mean
+
+            std = torch.sqrt(torch.var(inputs, dim=1, keepdim=True, unbiased=False) + 1e-5).detach()
+            inputs = inputs / std
+
+        h = self.encoder(inputs) 
+        hidden_size = h.size(1)
+        h = h.permute(0, 2, 3, 1).reshape(bs * c_num, length, hidden_size)  # (BC, L, H)
+
+        for i in range(length):
+            spks = self.net(h[:, i, :])
+
+        spks = spks.reshape(bs, c_num * hidden_size, -1)  # B, CH, Time Step
+
+        spks = spks[:, :, -1]  # aggregate over time dimension shape, (B, CH)
+        preds = self.fc(spks.view(bs, c_num, -1)).squeeze(-1) # B, O, C
+        preds = preds.permute(0, 2, 1).contiguous()
+        if self.args.normalize:
+            preds = preds * std + mean  # denormalize
+        aux = {'gate_l0': torch.tensor(0.0, device=preds.device)} # placeholder
+        return preds, aux
+
+    @property
+    def output_size(self):
+        return self.__output_size

model/TS_TCN.py

@@ -0,0 +1,1030 @@
+from typing import Optional, Callable
+
+import torch
+from torch import nn
+from torch.nn.utils import weight_norm
+import snntorch as snn
+from snntorch import surrogate
+from snntorch import utils
+import numpy as np
+import math
+import copy
+from spikingjelly.activation_based import surrogate, neuron
+from abc import abstractmethod
+import warnings
+
+
+
+surrogate.atan = lambda alpha=2.0: SG.apply
+
+
+class Chomp1d(nn.Module):
+    def __init__(self, chomp_size):
+        super().__init__()
+        self.chomp_size = chomp_size
+
+    def forward(self, x):
+        return x[:, :, : -self.chomp_size].contiguous()
+
+
+class Chomp2d(nn.Module):
+    def __init__(self, chomp_size):
+        super().__init__()
+        self.chomp_size = chomp_size
+
+    def forward(self, x):
+        return x[:, :, :, : -self.chomp_size].contiguous()
+
+
+
+class RepeatEncoder(nn.Module):
+    def __init__(self, output_size: int):
+        super().__init__()
+        self.out_size = output_size
+        self.lif = snn.Leaky(
+            beta=0.99, spike_grad=surrogate.atan(), init_hidden=True, output=False
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: batch, L, C
+        inputs = inputs.repeat(
+            tuple([self.out_size] + torch.ones(len(inputs.size()), dtype=int).tolist())
+        )  # out_size batch L C
+        inputs = inputs.permute(1, 0, 3, 2)  # batch out_size L C
+        spks = self.lif(inputs)
+        return spks
+
+
+class ConvEncoder(nn.Module):
+    def __init__(self, output_size: int, kernel_size: int = 3):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv2d(
+                in_channels=1,
+                out_channels=output_size,
+                kernel_size=(1, kernel_size),
+                stride=1,
+                padding=(0, kernel_size // 2),
+            ),
+            nn.BatchNorm2d(output_size),
+        )
+        self.lif = snn.Leaky(
+            beta=0.99,
+            spike_grad=surrogate.atan(alpha=2.0),
+            init_hidden=True,
+            output=False,
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: batch, L, C
+        inputs = inputs.permute(0, 2, 1).unsqueeze(1)  # batch, 1, C, L
+        enc = self.encoder(inputs)  # batch, output_size, C, L
+        spks = self.lif(enc)
+        return spks
+
+
+
+
+
+class DeltaEncoder(nn.Module):
+    def __init__(self, output_size: int):
+        super().__init__()
+        self.norm = nn.BatchNorm2d(1)
+        self.enc = nn.Linear(1, output_size)
+        self.lif = snn.Leaky(
+            beta=0.99, spike_grad=surrogate.atan(), init_hidden=True, output=False
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs: batch, L, C
+        delta = torch.zeros_like(inputs)
+        delta[:, 1:] = inputs[:, 1:, :] - inputs[:, :-1, :]
+        delta = delta.unsqueeze(1).permute(0, 1, 3, 2)  # batch, 1, C, L
+        delta = self.norm(delta)
+        delta = delta.permute(0, 2, 3, 1)  # batch, C, L, 1
+        enc = self.enc(delta)  # batch, C, L, output_size
+        enc = enc.permute(0, 3, 1, 2)  # batch, output_size, C, L
+        spks = self.lif(enc)
+        return spks
+
+
+SpikeEncoder = {
+    "snntorch": {
+        "repeat": RepeatEncoder,
+        "conv": ConvEncoder,
+        "delta": DeltaEncoder,
+    },
+    "spikingjelly": {
+        "repeat": RepeatEncoder,
+        "conv": ConvEncoder,
+        "delta": DeltaEncoder,
+    },
+}
+
+
+def generate_ones_and_minus_ones_matrix(rows, cols):
+    random_matrix = torch.randint(0, 2, (rows, cols))
+    binary_matrix = torch.where(
+        random_matrix == 0,
+        -1 * torch.ones_like(random_matrix),
+        torch.ones_like(random_matrix),
+    )
+    return binary_matrix.float()
+
+
+class RandomPE(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        pe_mode="concat",
+        num_pe_neuron=10,
+        neuron_pe_scale=1000.0,
+        dropout=0.1,
+        num_steps=4,
+    ):
+        super().__init__()
+        self.max_len = 5000  # different from windows
+        self.pe_mode = pe_mode
+        self.neuron_pe_scale = neuron_pe_scale
+        self.dropout = nn.Dropout(p=dropout)
+        if self.pe_mode == "concat":
+            self.num_pe_neuron = copy.deepcopy(num_pe_neuron)
+        elif self.pe_mode == "add":
+            self.num_pe_neuron = copy.deepcopy(d_model)
+        pe = generate_ones_and_minus_ones_matrix(
+            self.max_len, self.num_pe_neuron
+        )  # MaxL, Neur
+        pe = pe.unsqueeze(0).transpose(0, 1)  # MaxL, 1, Neur
+        print("pe.shape: ", pe.shape)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        # T, B, L, D
+        T, B, L, _ = x.shape
+        x = x.permute(1, 0, 2, 3)  # B, T, L, D
+        x = x.flatten(1, 2)  # B, TL, D
+        if self.pe_mode == "concat":
+            # tmp: TL, 1, Neur -> TL, B, Neur -> B, TL, Neur
+            tmp = self.pe[: x.size(-2), :].repeat(1, B, 1).transpose(0, 1)
+            x = torch.concat([x, tmp], dim=-1)
+            # print(x.shape) # B, TL, D'
+        elif self.pe_mode == "add":
+            # [B, TL, D] + [1, TL, Neur]
+            x = x + self.pe[: x.size(-2), :].transpose(0, 1)
+            # print(x.shape) # B, TL, D
+        x = x.transpose(0, 1)  # TL, B D
+        x = x.reshape(T, L, B, -1)  # T, L, B, D
+        x = x.permute(0, 2, 1, 3)  # T, B, L, D
+        return self.dropout(x)
+
+
+class NeuronPE(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        pe_mode="concat",
+        num_pe_neuron=10,
+        neuron_pe_scale=10000.0,
+        dropout=0.1,
+        num_steps=4,
+    ):
+        super().__init__()
+        self.max_len = 50000  # different from windows
+        self.pe_mode = pe_mode
+        self.neuron_pe_scale = neuron_pe_scale
+        self.dropout = nn.Dropout(p=dropout)
+        if self.pe_mode == "concat":
+            self.num_pe_neuron = copy.deepcopy(num_pe_neuron)
+        elif self.pe_mode == "add":
+            self.num_pe_neuron = copy.deepcopy(d_model)
+        pe = torch.zeros(self.max_len, self.num_pe_neuron)  # MaxL, Neur
+        position = torch.arange(0, self.max_len, dtype=torch.float).unsqueeze(
+            1
+        )  # MaxL, 1
+        div_term = torch.exp(
+            torch.arange(0, self.num_pe_neuron, 2).float()
+            * (-math.log(neuron_pe_scale) / self.num_pe_neuron)
+        )
+        div_term_single = torch.exp(
+            torch.arange(0, self.num_pe_neuron - 1, 2).float()
+            * (-math.log(neuron_pe_scale) / self.num_pe_neuron)
+        )
+        pe[:, 0::2] = torch.heaviside(
+            torch.sin(position * div_term) - 0.8, torch.tensor([1.0])
+        )
+        pe[:, 1::2] = torch.heaviside(
+            torch.cos(position * div_term_single) - 0.8, torch.tensor([1.0])
+        )
+        pe = pe.unsqueeze(0).transpose(0, 1)  # MaxL, 1, Neur
+        print("pe.shape: ", pe.shape)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        # T, B, L, D
+        T, B, L, _ = x.shape
+        x = x.permute(1, 0, 2, 3)  # B, T, L, D
+        x = x.flatten(1, 2)  # B, TL, D
+        if self.pe_mode == "concat":
+            # tmp: TL, 1, Neur -> TL, B, Neur -> B, TL, Neur
+            tmp = self.pe[: x.size(-2), :].repeat(1, B, 1).transpose(0, 1)
+            x = torch.concat([x, tmp], dim=-1)
+            # print(x.shape) # B, TL, D'
+        elif self.pe_mode == "add":
+            # [B, TL, D] + [1, TL, Neur]
+            # print(self.pe[:x.size(-2), :].shape)
+            x = x + self.pe[: x.size(-2), :].transpose(0, 1)
+            # print(x.shape) # B, TL, D
+        x = x.transpose(0, 1)  # TL, B D
+        x = x.reshape(T, L, B, -1)  # T, L, B, D
+        x = x.permute(0, 2, 1, 3)  # T, B, L, D
+        return self.dropout(x)
+
+
+class StaticPE(nn.Module):
+    r"""Inject some information about the relative or absolute position of the tokens
+        in the sequence. The positional encodings have the same dimension as
+        the embeddings, so that the two can be summed. Here, we use sine and cosine
+        functions of different frequencies.
+    .. math::
+        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
+        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
+        \text{where pos is the word position and i is the embed idx)"""
+
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        pe = torch.zeros(max_len, d_model)  # MaxL, D
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # MaxL, 1
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
+        )
+        div_term_single = torch.exp(
+            torch.arange(0, d_model - 1, 2).float() * (-math.log(10000.0) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term_single)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        # x: L, TB, D
+        x = x + self.pe[: x.size(0), :]
+        x = self.dropout(x)
+        return x
+
+
+class ConvPE(nn.Module):
+    def __init__(self, d_model, dropout=0.1, max_len=5000, num_steps=4):
+
+        super().__init__()
+        self.T = num_steps
+        self.rpe_conv = nn.Conv1d(
+            d_model, d_model, kernel_size=3, stride=1, padding=1, bias=False
+        )
+        self.rpe_bn = nn.BatchNorm1d(d_model)
+        self.rpe_lif = neuron.LIFNode(
+            step_mode="m",
+            detach_reset=True,
+            surrogate_function=surrogate.ATan(),
+            v_threshold=1.0,
+        )
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x):
+        # x: L, TB, D
+        L, TB, D = x.shape
+        x_feat = x.permute(1, 2, 0)  # TB, D, L
+        x_feat = self.rpe_conv(x_feat)  # TB, D, L
+        x_feat = (
+            self.rpe_bn(x_feat).reshape(self.T, int(TB / self.T), D, L).contiguous()
+        )  # T, B, D, L
+        x_feat = self.rpe_lif(x_feat)
+        x_feat = x_feat.flatten(0, 1)  # TB, D, L
+        x_feat = self.dropout(x_feat)  # TB, D, L
+        x_feat = x_feat.permute(2, 0, 1)  # L, TB, D
+        x = x + x_feat
+        return x
+
+
+class PositionEmbedding(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        pe_type: str,
+        max_len: int = 5000,
+        pe_mode: str = "add",
+        num_pe_neuron: int = 10,
+        neuron_pe_scale: float = 1000.0,
+        dropout=0.1,
+        num_steps=4,
+    ):
+        super().__init__()
+        self.emb_type = pe_type
+        if pe_type in ["learn", "none"]:
+            self.emb = nn.Embedding(max_len, input_size)
+        elif pe_type == "conv":
+            self.emb = ConvPE(
+                d_model=input_size,
+                max_len=max_len,
+                dropout=dropout,
+                num_steps=num_steps,
+            )
+        elif pe_type == "static":
+            self.emb = StaticPE(d_model=input_size, max_len=max_len, dropout=dropout)
+        elif pe_type == "neuron":
+            self.emb = NeuronPE(
+                d_model=input_size,
+                pe_mode=pe_mode,
+                num_pe_neuron=num_pe_neuron,
+                neuron_pe_scale=neuron_pe_scale,
+                dropout=dropout,
+                num_steps=num_steps,
+            )
+        elif pe_type == "random":
+            self.emb = RandomPE(
+                d_model=input_size,
+                pe_mode=pe_mode,
+                num_pe_neuron=num_pe_neuron,
+                neuron_pe_scale=neuron_pe_scale,
+                dropout=dropout,
+                num_steps=num_steps,
+            )
+        else:
+            raise ValueError("Unknown embedding type: {}".format(pe_type))
+
+    def forward(self, x):
+        if self.emb_type == "learn":
+            # T, B, L, D = x.shape # x: T, B, L, D
+            # x = x.flatten(0, 1) # TB, L, D
+            tmp = torch.arange(
+                end=x.size()[1], device=x.device
+            )  # [0,1,2,...,L-1], shape: L
+            embedding = self.emb(tmp)  # shape: L, D
+            embedding = embedding.repeat([x.size()[0], 1, 1])  # TB, L, D'
+            x = x + embedding
+            # x = x.reshape(T, B, L, -1)
+        elif self.emb_type in ["static", "conv"]:
+            T, B, L, _ = x.shape  # x: T, B, L, D
+            x = x.flatten(0, 1)  # TB, L, D
+            x = self.emb(x.transpose(0, 1)).transpose(0, 1)  # x: TB, L, D'
+            x = x.reshape(T, B, L, -1)
+        elif self.emb_type in ["neuron", "random"]:
+            T, B, L, _ = x.shape  # x: T, B, L, D
+            # T, B, L, D
+            x = self.emb(x)
+            x = x.reshape(T, B, L, -1)
+        return x  # T, B, L, D'
+
+
+@torch.jit.script
+def heaviside(x: torch.Tensor):
+    return (x >= 0).to(x)
+
+@torch.jit.script
+def atan_backward(grad_output: torch.Tensor, x: torch.Tensor, alpha: float):
+
+    return alpha / 2 / (1 + (math.pi / 2 * alpha * x).pow_(2)) * grad_output, None
+             
+
+class SG(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, alpha=2.0):
+        if x.requires_grad:
+            ctx.save_for_backward(x)
+            ctx.alpha = alpha
+        return heaviside(x)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return atan_backward(grad_output, ctx.saved_tensors[0], ctx.alpha)
+
+
+class MemoryModule(nn.Module):
+    def __init__(self):
+        """
+        * :ref:`API in English <MemoryModule.__init__-en>`
+
+        .. _MemoryModule.__init__-cn:
+
+        ``MemoryModule`` 是SpikingJelly中所有有状态（记忆）模块的基类。
+
+        * :ref:`中文API <MemoryModule.__init__-cn>`
+
+        .. _MemoryModule.__init__-en:
+
+        ``MemoryModule`` is the base class of all stateful modules in SpikingJelly.
+
+        """
+        super().__init__()
+        self._memories = {}
+        self._memories_rv = {}
+
+    def register_memory(self, name: str, value):
+        """
+        * :ref:`API in English <MemoryModule.register_memory-en>`
+
+        .. _MemoryModule.register_memory-cn:
+
+        :param name: 变量的名字
+        :type name: str
+        :param value: 变量的值
+        :type value: any
+
+        将变量存入用于保存有状态变量（例如脉冲神经元的膜电位）的字典中。这个变量的重置值会被设置为 ``value``。每次调用 ``self.reset()``
+        函数后， ``self.name`` 都会被重置为 ``value``。
+
+        * :ref:`中文API <MemoryModule.register_memory-cn>`
+
+        .. _MemoryModule.register_memory-en:
+
+        :param name: variable's name
+        :type name: str
+        :param value: variable's value
+        :type value: any
+
+        Register the variable to memory dict, which saves stateful variables (e.g., the membrane potential of a
+        spiking neuron). The reset value of this variable will be ``value``. ``self.name`` will be set to ``value`` after
+        each calling of ``self.reset()``.
+
+        """
+        assert not hasattr(self, name), f'{name} has been set as a member variable!'
+        self._memories[name] = value
+        self.set_reset_value(name, value)
+
+    def reset(self):
+        """
+        * :ref:`API in English <MemoryModule.reset-en>`
+
+        .. _MemoryModule.reset-cn:
+
+        重置所有有状态变量为默认值。
+
+        * :ref:`中文API <MemoryModule.reset-cn>`
+
+        .. _MemoryModule.reset-en:
+
+        Reset all stateful variables to their default values.
+        """
+        for key in self._memories.keys():
+            self._memories[key] = copy.deepcopy(self._memories_rv[key])
+
+    def set_reset_value(self, name: str, value):
+        self._memories_rv[name] = copy.deepcopy(value)
+
+    def __getattr__(self, name: str):
+        if '_memories' in self.__dict__:
+            memories = self.__dict__['_memories']
+            if name in memories:
+                return memories[name]
+
+        return super().__getattr__(name)
+
+    def __setattr__(self, name: str, value) -> None:
+        _memories = self.__dict__.get('_memories')
+        if _memories is not None and name in _memories:
+            _memories[name] = value
+        else:
+            super().__setattr__(name, value)
+
+    def __delattr__(self, name):
+        if name in self._memories:
+            del self._memories[name]
+            del self._memories_rv[name]
+        else:
+            return super().__delattr__(name)
+
+    def __dir__(self):
+        module_attrs = dir(self.__class__)
+        attrs = list(self.__dict__.keys())
+        parameters = list(self._parameters.keys())
+        modules = list(self._modules.keys())
+        buffers = list(self._buffers.keys())
+        memories = list(self._memories.keys())
+        keys = module_attrs + attrs + parameters + modules + buffers + memories
+
+        # Eliminate attrs that are not legal Python variable names
+        keys = [key for key in keys if not key[0].isdigit()]
+
+        return sorted(keys)
+
+    def memories(self):
+        """
+        * :ref:`API in English <MemoryModule.memories-en>`
+
+        .. _MemoryModule.memories-cn:
+
+        :return: 返回一个所有状态变量的迭代器
+        :rtype: Iterator
+
+        * :ref:`中文API <MemoryModule.memories-cn>`
+
+        .. _MemoryModule.memories-en:
+
+        :return: an iterator over all stateful variables
+        :rtype: Iterator
+        """
+        for name, value in self._memories.items():
+            yield value
+
+    def named_memories(self):
+        """
+        * :ref:`API in English <MemoryModule.named_memories-en>`
+
+        .. _MemoryModule.named_memories-cn:
+
+        :return: 返回一个所有状态变量及其名称的迭代器
+        :rtype: Iterator
+
+        * :ref:`中文API <MemoryModule.named_memories-cn>`
+
+        .. _MemoryModule.named_memories-en:
+
+        :return: an iterator over all stateful variables and their names
+        :rtype: Iterator
+        """
+
+        for name, value in self._memories.items():
+            yield name, value
+
+    def detach(self):
+        """
+        * :ref:`API in English <MemoryModule.detach-en>`
+
+        .. _MemoryModule.detach-cn:
+
+        从计算图中分离所有有状态变量。
+
+        .. tip::
+
+            可以使用这个函数实现TBPTT(Truncated Back Propagation Through Time)。
+
+
+        * :ref:`中文API <MemoryModule.detach-cn>`
+
+        .. _MemoryModule.detach-en:
+
+        Detach all stateful variables.
+
+        .. admonition:: Tip
+            :class: tip
+
+            We can use this function to implement TBPTT(Truncated Back Propagation Through Time).
+
+        """
+
+        for key in self._memories.keys():
+            if isinstance(self._memories[key], torch.Tensor):
+                self._memories[key].detach_()
+
+    def _apply(self, fn):
+        for key, value in self._memories.items():
+            if isinstance(value, torch.Tensor):
+                self._memories[key] = fn(value)
+        return super()._apply(fn)
+
+    def _replicate_for_data_parallel(self):
+        replica = super()._replicate_for_data_parallel()
+        replica._memories = self._memories.copy()
+        return replica
+
+
+class StepModule:
+    def supported_step_mode(self):
+        """
+        * :ref:`API in English <StepModule.supported_step_mode-en>`
+        .. _StepModule.supported_step_mode-cn:
+        :return: 包含支持的后端的tuple
+        :rtype: tuple[str]
+        返回此模块支持的步进模式。
+        * :ref:`中文 API <StepModule.supported_step_mode-cn>`
+        .. _StepModule.supported_step_mode-en:
+        :return: a tuple that contains the supported backends
+        :rtype: tuple[str]
+        """
+        return ('s', 'm')
+
+    @property
+    def step_mode(self):
+        """
+        * :ref:`API in English <StepModule.step_mode-en>`
+        .. _StepModule.step_mode-cn:
+        :return: 模块当前使用的步进模式
+        :rtype: str
+        * :ref:`中文 API <StepModule.step_mode-cn>`
+        .. _StepModule.step_mode-en:
+        :return: the current step mode of this module
+        :rtype: str
+        """
+        return self._step_mode
+
+    @step_mode.setter
+    def step_mode(self, value: str):
+        """
+        * :ref:`API in English <StepModule.step_mode-setter-en>`
+        .. _StepModule.step_mode-setter-cn:
+        :param value: 步进模式
+        :type value: str
+        将本模块的步进模式设置为 ``value``
+        * :ref:`中文 API <StepModule.step_mode-setter-cn>`
+        .. _StepModule.step_mode-setter-en:
+        :param value: the step mode
+        :type value: str
+        Set the step mode of this module to be ``value``
+        """
+        if value not in self.supported_step_mode():
+            raise ValueError(f'step_mode can only be {self.supported_step_mode()}, but got "{value}"!')
+        self._step_mode = value
+
+
+
+class BaseNode(MemoryModule):
+    def __init__(self,
+                 v_threshold: float = 1.,
+                 v_reset: float = 0.,
+                 surrogate_function: Callable = None,
+                 detach_reset: bool = False,
+                 step_mode='s', backend='torch',
+                 store_v_seq: bool = True):
+
+        assert isinstance(v_reset, float) or v_reset is None
+        assert isinstance(v_threshold, float)
+        assert isinstance(detach_reset, bool)
+        super().__init__()
+
+        if v_reset is None:
+            self.register_memory('v', 0.)
+            self.register_memory('v_s', 0.)
+        else:
+            self.register_memory('v', v_reset)
+
+        self.v_threshold = v_threshold
+
+        self.v_reset = v_reset
+        self.detach_reset = detach_reset
+        self.surrogate_function = surrogate_function
+
+        self.step_mode = step_mode
+        self.backend = backend
+
+        self.store_v_seq = store_v_seq
+
+
+        self.alpha_s = torch.nn.Parameter(torch.tensor(0.5, dtype=torch.float)) 
+        self.alpha_l = torch.nn.Parameter(torch.tensor(0.5, dtype=torch.float))
+
+    @property
+    def store_v_seq(self):
+        return self._store_v_seq
+
+    @store_v_seq.setter
+    def store_v_seq(self, value: bool):
+        self._store_v_seq = value
+        if value:
+            if not hasattr(self, 'v_seq'):
+                self.register_memory('v_seq', None)
+
+    @staticmethod
+    @torch.jit.script
+    def jit_hard_reset(v: torch.Tensor, spike: torch.Tensor, v_reset: float):
+        v = (1. - spike) * v + spike * v_reset
+
+        return v
+
+    @staticmethod
+    @torch.jit.script
+    def jit_soft_reset(v: torch.Tensor, spike: torch.Tensor, v_threshold: float):
+        v = v - spike * v_threshold
+        return v
+
+
+    @abstractmethod
+    def neuronal_charge(self, x: torch.Tensor):
+        raise NotImplementedError
+
+    def neuronal_fire(self):
+        return self.surrogate_function(self.v - self.v_threshold, 2.0)
+
+    def sl_neuronal_fire(self):
+        s_s = self.surrogate_function(self.v - self.v_threshold, 2.0)
+        s_l = self.surrogate_function(self.v_s - self.v_threshold,  2.0)
+        return s_s, s_l
+
+    def extra_repr(self):
+        return f'v_threshold={self.v_threshold}, v_reset={self.v_reset}, detach_reset={self.detach_reset}, step_mode={self.step_mode}, backend={self.backend}'
+
+    def single_step_forward(self, x: torch.Tensor):
+        self.v_float_to_tensor(x)
+        self.neuronal_charge(x)
+        s_s, s_l = self.sl_neuronal_fire()
+        spike = self.alpha_s * s_s + self.alpha_l * s_l
+        self.neuronal_reset(s_s, s_l)
+        
+        return spike
+
+    def multi_step_forward(self, x_seq: torch.Tensor):
+
+        T = x_seq.shape[-1]
+        y_seq = []
+        if self.store_v_seq:
+            v_seq = []
+        for t in range(T):
+            y = self.single_step_forward(x_seq[:, t])
+            y_seq.append(y)
+            if self.store_v_seq:
+                v_seq.append(self.v)
+        if self.store_v_seq:
+            self.v_seq = torch.stack(v_seq)
+
+        outputs = torch.stack(y_seq, dim=0).permute(1, 0)
+
+        return outputs
+
+    def v_float_to_tensor(self, x: torch.Tensor):
+        if isinstance(self.v, float):
+            v_init = self.v
+            self.v = torch.full_like(x.data, v_init)
+
+
+class TSLIFNode(BaseNode):
+    def __init__(self,
+                 v_threshold=1.0,
+                 v_reset=0.,
+                 surrogate_function: Callable = None,
+                 detach_reset=False,
+                 hard_reset=False,
+                 step_mode='s',
+                 k=2,
+                 decay_factor: torch.Tensor = torch.tensor([0.8, 0.2, 0.3, 0.7], dtype=torch.float),
+                 gamma: float = 0.5):
+        super(TSLIFNode, self).__init__(v_threshold, v_reset, surrogate_function, detach_reset, step_mode)
+        self.k = k
+        for i in range(1, self.k + 1):
+            self.register_memory('v' + str(i), 0.)
+        self.names = self._memories
+        self.hard_reset = hard_reset
+        self.gamma = gamma
+        self.decay_factor = torch.nn.Parameter(decay_factor)
+        self.kk = torch.nn.Parameter(torch.tensor([0.8], dtype=torch.float))
+        self.yy = torch.nn.Parameter(torch.tensor([0.1], dtype=torch.float))
+
+    @property
+    def supported_backends(self):
+        if self.step_mode == 's':
+            return ('torch',)
+        elif self.step_mode == 'm':
+            return ('torch', 'cupy')
+        else:
+            raise ValueError(self.step_mode)
+
+    def neuronal_charge(self, x: torch.Tensor):
+        self.names['v1'] = self.decay_factor[0] * self.names['v1'] + self.decay_factor[1] * x - self.yy * self.names['v2']
+        self.names['v2'] = self.decay_factor[2] * self.names['v2'] + self.decay_factor[3] * x - self.kk * self.names['v1']
+        self.v = self.names['v2']
+        self.v_s = self.names['v1']
+
+    def neuronal_reset(self, spike_s, spike_l):
+        if not self.hard_reset:
+            self.names['v1'] = self.jit_soft_reset(self.names['v1'], spike_l, self.gamma)
+            self.names['v2'] = self.jit_soft_reset(self.names['v2'], spike_s, self.v_threshold)
+        else:
+            for i in range(2, self.k + 1):
+                self.names['v' + str(i)] = self.jit_hard_reset(self.names['v' + str(i)], spike_s, self.v_reset)
+
+    def forward(self, x: torch.Tensor):
+        return super().single_step_forward(x)
+    def extra_repr(self):
+        return f"v_threshold={self.v_threshold}, v_reset={self.v_reset}, detach_reset={self.detach_reset}, " \
+               f"hard_reset={self.hard_reset}, " \
+               f"gamma={self.gamma}, k={self.k}, step_mode={self.step_mode}, backend={self.backend}"
+
+
+
+
+
+class SpikeTemporalBlock(nn.Module):
+    def __init__(
+        self,
+        n_inputs,
+        n_outputs,
+        kernel_size,
+        stride,
+        dilation,
+        padding,
+        num_steps=4,
+    ):
+        super().__init__()
+        self.num_steps = num_steps
+        self.conv1 = weight_norm(
+            nn.Conv2d(
+                n_inputs,
+                n_outputs,
+                (1, kernel_size),
+                stride=stride,
+                padding=(0, padding),
+                dilation=(1, dilation),
+            )
+        )
+        self.bn1 = nn.BatchNorm2d(n_outputs)
+        self.chomp1 = Chomp2d(padding)
+
+        self.tslif1 = TSLIFNode(
+            surrogate_function =SG.apply,
+        )
+
+        self.conv2 = weight_norm(
+            nn.Conv2d(
+                n_outputs,
+                n_outputs,
+                (1, kernel_size),
+                stride=stride,
+                padding=(0, padding),
+                dilation=(1, dilation),
+            )
+        )
+        self.bn2 = nn.BatchNorm2d(n_outputs)
+        self.chomp2 = Chomp2d(padding)
+
+        self.tslif2 = TSLIFNode(
+            surrogate_function =SG.apply,
+        )
+
+        self.downsample = (
+            nn.Conv2d(n_inputs, n_outputs, (1, 1)) if n_inputs != n_outputs else None
+        )
+
+        self.tslif = TSLIFNode(
+            surrogate_function =SG.apply,
+        )
+
+    def init_weights(self):
+        self.conv1.weight.data.normal_(0, 0.01)
+        self.conv2.weight.data.normal_(0, 0.01)
+        if self.downsample is not None:
+            self.downsample.weight.data.normal_(0, 0.01)
+
+    def forward(self, x):
+        # out1: 24, 16, 361, 168
+        
+        out1 = self.chomp1(self.bn1(self.conv1(x)))
+        spk_rec1 = []
+        for _ in range(self.num_steps):
+            spk = self.tslif1(out1)
+            spk_rec1.append(spk)
+
+        spks1 = torch.stack(spk_rec1, dim=-1)  # spks1: B, H, C, L, T
+        spks1 = spks1.mean(-1)  # spks1: B, H, C, L
+
+        out2 = self.chomp2(self.bn2(self.conv2(spks1)))
+        spk_rec2 = []
+        for _ in range(self.num_steps):
+            # spk: 24, 16, 361, 168
+            spk = self.tslif2(out2)
+            spk_rec2.append(spk)
+
+        spks2 = torch.stack(spk_rec2, dim=-1)  # spks2: B, H, C, L, T
+        spks2 = spks2.mean(-1)  # spks2: B, H, C, L
+
+        if torch.isnan(spks2).any() or torch.isinf(spks2).any():
+            print("illegal value in TemporalBlock2D")
+
+        if self.downsample is None:
+            res = x
+        else:
+            res = self.downsample(x)
+
+        spk_rec3 = []
+        for _ in range(self.num_steps):
+
+            spk = self.tslif(spks2 + res)
+            spk_rec3.append(spk)
+
+
+        res = torch.stack(spk_rec3, dim=-1)  # res: B, H, C, L, T
+
+        res = res.mean(-1)
+
+        return res
+
+
+
+
+
+class TSTCN(nn.Module):
+    def __init__(
+        self,
+        args,
+        num_levels: int = 3,
+        channel: int = 16,
+        dilation: int = 2,
+        stride: int = 1,
+        kernel_size: int = 2,
+        dropout: float = 0.2,
+        max_length: int = 100,
+        encoder_type: str = "conv",
+        pe_type: str = "neuron",
+        pe_mode: str = "concat",
+        num_pe_neuron: int = 40,
+        neuron_pe_scale: float = 1000.0,
+    ):
+        super().__init__()
+        
+        self.hidden_size   = args.hidden_size  
+        self.num_steps = args.T    
+        self.input_size = args.feature_size  
+        self.feature_size = args.feature_size
+        self.pre_length = args.pre_length      
+        self.num_levels = args.blocks
+        self.pe_type = pe_type
+        self.pe_mode = pe_mode
+        self.num_pe_neuron = num_pe_neuron
+        self.kernel_size = args.kernel_size
+        self.args = args
+        
+
+
+        self._snn_backend = "snntorch"
+        self.encoder = SpikeEncoder[self._snn_backend][encoder_type](self.hidden_size)
+        
+
+        self.pe = PositionEmbedding(
+            pe_type=pe_type,
+            pe_mode=pe_mode,
+            neuron_pe_scale=neuron_pe_scale,
+            input_size=self.input_size,
+            max_len=max_length,
+            num_pe_neuron=self.num_pe_neuron,
+            dropout=0.1,
+            num_steps=self.num_steps,
+        )
+
+
+
+        layers = []
+        num_channels = [channel] * self.num_levels
+        num_channels.append(1)
+        for i in range(self.num_levels + 1):
+            dilation_size = dilation**i
+            in_channels = self.hidden_size if i == 0 else num_channels[i - 1]
+            out_channels = num_channels[i]
+            layers += [
+                SpikeTemporalBlock(
+                    in_channels,
+                    out_channels,
+                    self.kernel_size,
+                    stride=stride,
+                    dilation=dilation_size,
+                    padding=(self.kernel_size - 1) * dilation_size,
+                    num_steps=self.num_steps,
+                )
+            ]
+        
+
+        
+        self.network = nn.Sequential(*layers)
+
+        if (self.pe_type == "neuron" and self.pe_mode == "concat") or (
+            self.pe_type == "random" and self.pe_mode == "concat"
+        ):
+            self.__output_size = self.feature_size + num_pe_neuron
+        else:
+            self.__output_size = args.seq_length
+
+        self.fc1 = nn.Linear(self.__output_size, args.feature_size)
+        self.fc2 = nn.Linear(args.seq_length, self.pre_length)
+        self.to('cuda:0') 
+    
+    def forward(self, inputs: torch.Tensor):
+        utils.reset(self.encoder)
+
+        if self.args.normalize:
+
+            mean = inputs.mean(dim=1, keepdim=True).detach() # shape [B, 1, D]
+            inputs = inputs - mean
+
+            std = torch.sqrt(torch.var(inputs, dim=1, keepdim=True, unbiased=False) + 1e-5).detach()
+            inputs = inputs / std
+
+        inputs = self.encoder(inputs)  # B, H, C, L
+        # inputs: 24, 64, 321, 168
+        
+    
+        if self.pe_type != "none":
+            inputs = self.pe(inputs.permute(1, 0, 3, 2)).permute(1, 0, 3, 2)
+        
+        spks = self.network(inputs)
+        spks = spks.squeeze(1)  # B, C', L
+        preds = self.fc1(spks.permute(0, 2, 1))  # B, L, C
+        preds = self.fc2(preds.permute(0, 2, 1))  # B, C', L
+        preds = preds.permute(0, 2, 1).contiguous()
+        if self.args.normalize:
+            preds = preds * std + mean  # denormalize
+        
+        
+        # Create auxiliary output
+        aux = {'gate_l0': torch.tensor(0.0, device=preds.device)} # placeholder
+        
+        return preds, aux
+    
+    @property
+    def output_size(self):
+        return self.__output_size

model/iSpikformer.py

@@ -0,0 +1,129 @@
+import torch
+from torch import nn
+from spikingjelly.clock_driven.neuron import MultiStepLIFNode
+
+class SPE(nn.Module):
+    def __init__(self, input_len, patch_num, patch_dim, T, tau, D):
+        super().__init__()
+        self.patch_projector = nn.Linear(input_len // patch_num, patch_dim)
+        self.bn = nn.BatchNorm2d(patch_dim)
+        self.encoder_lif = MultiStepLIFNode(tau=tau, detach_reset=False, backend='torch')
+
+        self.D = D
+        self.T = T
+        self.patch_dim = patch_dim
+        self.patch_num = patch_num
+
+    def forward(self, x):
+        B, L, D = x.shape
+
+        x = x.view(B, self.patch_num, L // self.patch_num, D).contiguous()
+        x = x.transpose(-1, -2).contiguous()
+        x = self.patch_projector(x)
+        x = x.repeat(self.T, 1, 1, 1, 1)
+        x = x.permute(0, 1, 4, 2, 3).contiguous()
+        x = x.flatten(0, 1)
+        x = self.bn(x)
+        x = x.view(self.T, B, self.patch_dim, self.patch_num, D)
+        x = self.encoder_lif(x)
+
+        return x
+
+class iSSA(nn.Module):
+    def __init__(self, patch_num, D, patch_dim, tau, alpha):
+        super().__init__()
+        self.lin1 = nn.Linear(patch_num, patch_num)
+        self.lin2 = nn.Linear(patch_num, patch_num)
+        self.lin3 = nn.Linear(patch_num, patch_num)
+
+        self.lif1 = MultiStepLIFNode(tau=tau, detach_reset=True, backend='torch')
+        self.lif2 = MultiStepLIFNode(tau=tau, detach_reset=True, backend='torch')
+        self.lif3 = MultiStepLIFNode(tau=tau, detach_reset=True, backend='torch')
+        self.lif4 = MultiStepLIFNode(tau=tau, detach_reset=True, backend='torch')
+
+        self.b1 = nn.BatchNorm2d(patch_dim)
+        self.b2 = nn.BatchNorm2d(patch_dim)
+        self.b3 = nn.BatchNorm2d(patch_dim)
+        self.b4 = nn.BatchNorm2d(patch_dim)
+     
+    def forward(self, x):
+        res_x = x
+        T, B, pd, pn, D = x.shape
+
+
+        x = x.transpose(-1, -2).contiguous()
+        q = self.lin1(x).flatten(0, 1)
+        k = self.lin2(x).flatten(0, 1)
+        v = self.lin3(x).flatten(0, 1)
+
+        q = self.b1(q)
+        k = self.b2(k)
+        v = self.b3(v)
+
+        q = q.view(T, B, pd, D, -1)
+        k = k.view(T, B, pd, D, -1)
+        v = v.view(T, B, pd, D, -1)
+
+        q = self.lif1(q)
+        k = self.lif2(k).transpose(-1, -2).contiguous()
+        v = self.lif3(v)
+
+        attn = q @ k
+        attn = attn @ v
+        attn = attn.flatten(0, 1)
+        attn = self.b4(attn)
+        attn = attn.view(T, B, pd, D, pn)
+        attn = self.lif4(attn)
+        attn = attn.transpose(-1, -2).contiguous()
+
+        return attn
+
+class iSpikformer(nn.Module):
+    def __init__(self, args, input_len, patch_num, patch_dim, T, blocks, D, pred_len, tau, alpha, hidden_dim):
+        super().__init__()
+        self.emb = SPE(input_len, patch_num, patch_dim, T, tau, D)
+        self.args = args
+        self.attn = nn.ModuleList()
+        for i in range(blocks):
+            self.attn.append(iSSA(patch_num, D, patch_dim, tau, alpha))
+
+        self.dense1 = nn.Linear(patch_num*patch_dim, hidden_dim)
+        self.dense2 = nn.Linear(hidden_dim, pred_len)
+        self.bn = nn.BatchNorm1d(D)
+        self.activ = MultiStepLIFNode(tau=tau, detach_reset=True, backend='torch')
+        self.to('cuda:0') 
+
+
+    def forward(self, x):
+        if self.args.normalize:
+            mean = x.mean(dim=1, keepdim=True).detach()
+            x = x - mean
+
+            std = torch.sqrt(torch.var(x, dim=1, keepdim=True, unbiased=False) + 1e-5).detach()
+            x = x / std
+
+        x = self.emb(x)
+        T, B, pd, pn, D = x.shape
+
+        for i in range(len(self.attn)):
+            x = self.attn[i](x)
+        x = x.permute(0, 1, 4, 2, 3).contiguous()
+        x = x.flatten(-2, -1)
+        x = self.dense1(x)
+        x = x.flatten(0, 1)
+        x = self.bn(x)
+        x = self.activ(x)
+        x = self.dense2(x)
+        x = x.transpose(-1, -2).contiguous()
+        x = x.view(T, B, -1, D)
+
+        if self.args.normalize:
+            x = x * std
+            x = x + mean.repeat(T, 1, 1, 1)
+
+        aux = {
+            'gate_l0': torch.tensor(0.0, device=x.device)  # placeholder
+        }
+        
+        return x.mean(dim=0), aux
+    

requirements.txt

@@ -0,0 +1,6 @@
+numpy
+pandas
+torch
+scikit-learn
+snntorch
+spikingjelly

scripts/ecl.sh

@@ -0,0 +1,232 @@
+#!/bin/bash
+
+
+
+
+
+
+if [ ! -d "./logs" ]; then
+    mkdir ./logs
+fi
+
+if [ ! -d "./logs/LongForecasting" ]; then
+    mkdir ./logs/LongForecasting
+fi
+
+
+
+
+
+
+
+
+python train.py \
+    --model FGN \
+    --data electricity \
+    --feature_size 370\
+    --embed_size 128 \
+    --hidden_size 256 \
+    --batch_size 16 \
+    --train_ratio 0.7 \
+    --val_ratio 0.2 \
+    --seq_length 12 \
+    --pre_length 12 \
+    --train_epochs 100 \
+    --learning_rate 0.00001 \
+    --device cuda:0 >logs/LongForecasting/ECL_FGN.log
+
+
+python train.py \
+    --model SpikF \
+    --data electricity \
+    --feature_size 370\
+    --embed_size 128 \
+    --hidden_size 256 \
+    --batch_size 16 \
+    --train_ratio 0.7 \
+    --val_ratio 0.2 \
+    --seq_length 12 \
+    --pre_length 12 \
+    --train_epochs 100 \
+    --learning_rate 0.00001 \
+    --T 16 \
+    --blocks 2\
+    --device cuda:0 >logs/LongForecasting/ECL_SpikF.log
+
+
+python train.py \
+    --model iSpikformer \
+    --data electricity \
+    --feature_size 370\
+    --embed_size 128 \
+    --hidden_size 256 \
+    --batch_size 16 \
+    --train_ratio 0.7 \
+    --val_ratio 0.2 \
+    --seq_length 12 \
+    --pre_length 12 \
+    --train_epochs 100 \
+    --learning_rate 0.00001 \
+    --blocks 2 \
+    --device cuda:0 >logs/LongForecasting/ECL_iSpikformer.log
+
+
+
+python train.py \
+    --model SpikF_GO \
+    --data electricity \
+    --feature_size 370\
+    --embed_size 128 \
+    --hidden_size 256 \
+    --batch_size 16 \
+    --train_ratio 0.7 \
+    --val_ratio 0.2 \
+    --seq_length 12 \
+    --pre_length 12 \
+    --train_epochs 100 \
+    --learning_rate 0.00001 \
+    --energy_loss True \
+    --device cuda:0 >logs/LongForecasting/ECL_SpikFGO.log
+
+
+python train.py \
+    --model SpikF_GO_CPG \
+    --data electricity \
+    --feature_size 370\
+    --embed_size 128 \
+    --hidden_size 256 \
+    --batch_size 16 \
+    --train_ratio 0.7 \
+    --val_ratio 0.2 \
+    --seq_length 12 \
+    --pre_length 12 \
+    --train_epochs 100 \
+    --learning_rate 0.00001 \
+    --energy_loss True \
+    --device cuda:0 >logs/LongForecasting/ECL_SpikFGOCPG.log
+
+
+python train.py \
+    --model SpikeRNN_CPG \
+    --data electricity \
+    --feature_size 370\
+    --embed_size 128 \
+    --hidden_size 128\
+    --batch_size 16 \
+    --train_ratio 0.7 \
+    --val_ratio 0.2 \
+    --seq_length 12 \
+    --pre_length 12 \
+    --train_epochs 100 \
+    --learning_rate 0.00001 \
+    --blocks 2 \
+    --device cuda:0 >logs/LongForecasting/ECL_SpikeRNNCPG.log
+
+
+
+
+python train.py \
+    --model SpikeGRU \
+    --data electricity \
+    --feature_size 370\
+    --embed_size 128 \
+    --hidden_size 64 \
+    --batch_size 16 \
+    --train_ratio 0.7 \
+    --val_ratio 0.2 \
+    --seq_length 12 \
+    --pre_length 12 \
+    --train_epochs 100 \
+    --learning_rate 0.00001 \
+    --device cuda:0 >logs/LongForecasting/ECL_SpikeGRU.log
+
+
+
+
+python train.py \
+    --model SpikeTCN_CPG \
+    --data electricity \
+    --feature_size 370\
+    --embed_size 128 \
+    --hidden_size 64\
+    --batch_size 16 \
+    --train_ratio 0.7 \
+    --val_ratio 0.2 \
+    --seq_length 12 \
+    --pre_length 12 \
+    --train_epochs 100 \
+    --learning_rate 0.00001 \
+    --blocks 3\
+    --device cuda:0 >logs/LongForecasting/ECL_SpikeTCNCPG.log
+
+
+
+
+python train.py \
+    --model Spikformer_CPG \
+    --data electricity \
+    --feature_size 370\
+    --embed_size 128 \
+    --hidden_size 128\
+    --batch_size 16 \
+    --train_ratio 0.7 \
+    --val_ratio 0.2 \
+    --seq_length 12 \
+    --pre_length 12 \
+    --train_epochs 100 \
+    --learning_rate 0.00001 \
+    --blocks 2 \
+    --device cuda:0 >logs/LongForecasting/ECL_SpikformerCPG.log
+
+
+
+python train.py \
+    --model TSTCN \
+    --data electricity \
+    --feature_size 370\
+    --embed_size 128 \
+    --hidden_size 64 \
+    --batch_size 16 \
+    --train_ratio 0.7 \
+    --val_ratio 0.2 \
+    --seq_length 12 \
+    --pre_length 12 \
+    --train_epochs 100 \
+    --learning_rate 0.00001 \
+    --kernel_size 3\
+    --blocks 3 \
+    --device cuda:0 >logs/LongForecasting/ECL_TSTCN.log
+
+
+python train.py \
+    --model TSGRU \
+    --data electricity \
+    --feature_size 370\
+    --embed_size 128 \
+    --hidden_size 64 \
+    --batch_size 16 \
+    --train_ratio 0.7 \
+    --val_ratio 0.2 \
+    --seq_length 12 \
+    --pre_length 12 \
+    --train_epochs 100 \
+    --learning_rate 0.00001 \
+    --device cuda:0 >logs/LongForecasting/ECL_TSGRU.log
+
+
+
+
+python train.py \
+    --model TSFormer \
+    --data electricity \
+    --feature_size 370\
+    --embed_size 128 \
+    --hidden_size 64 \
+    --batch_size 16 \
+    --train_ratio 0.7 \
+    --val_ratio 0.2 \
+    --seq_length 12 \
+    --pre_length 12 \
+    --train_epochs 100 \
+    --learning_rate 0.00001 \
+    --device cuda:0 >logs/LongForecasting/ECL_TSFormer.log

train.py

@@ -0,0 +1,545 @@
+import argparse
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import snntorch as snn
+import time
+import os
+import numpy as np
+import warnings
+from spikingjelly.clock_driven import functional
+
+from data.data_loader import (
+    Dataset_ECG, Dataset_Dhfm, Dataset_Solar, Dataset_Wiki, Dataset_PEMS_BAY
+)
+from utils.utils import save_model_ts, load_model_ts, evaluate
+
+from model.FourierGNN import FGN
+from model.SpikF import SpikF
+from model.iSpikformer import iSpikformer
+from model.SpikF_GO import SpikF_GO
+from model.SpikF_GO_CPG import SpikF_GO_CPG
+from model.TS_GRU import TSGRU
+from model.TS_TCN import TSTCN
+from model.TS_Former import TSFormer
+from model.SpikeGRU import SpikeGRU
+from model.Spikformer_CPG import Spikformer_CPG
+from model.SpikeRNN_CPG import SpikeRNN_CPG
+from model.SpikeTCN_CPG import SpikeTCN_CPG
+from model.TS_TCN import TSLIFNode
+
+
+def remove(model):
+    """Reset states of spiking neurons with warning suppression"""
+    if model is None:
+        return
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", message=".*not base.MemoryModule.*")
+        if hasattr(model, '__iter__'):
+            for m in model:
+                if hasattr(m, 'reset'):
+                    m.reset()
+                elif hasattr(m, 'v'):
+                    m.v = 0.0
+        elif hasattr(model, 'reset'):
+            model.reset()
+        elif hasattr(model, 'v'):
+            model.v = 0.0
+
+
+def reset_states(model):
+    """Reset states of all spiking neurons (TSLIFNode, Leaky, etc.) with warning suppression."""
+    if model is None:
+        return
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", message=".*not base.MemoryModule.*")
+        if hasattr(model, '__iter__'):
+            for m in model:
+                reset_states(m)
+        elif hasattr(model, 'modules'):
+            for module in model.modules():
+                if isinstance(module, (snn.Leaky, TSLIFNode)):
+                    try:
+                        module.reset()
+                    except Exception:
+                        if hasattr(module, 'v'):
+                            module.v = 0.0
+        elif hasattr(model, 'reset'):
+            model.reset()
+        elif hasattr(model, 'v'):
+            model.v = 0.0
+
+
+def _inverse_if_possible(arr: np.ndarray, scaler):
+    """
+    Inverse-transform arr of shape (..., D) using scaler fitted on train.
+    If scaler is None, returns arr unchanged.
+    """
+    if scaler is None:
+        return arr
+    if not hasattr(scaler, "inverse_transform"):
+        return arr
+
+    if arr.ndim < 2:
+        return arr
+
+    D = arr.shape[-1]
+    flat = arr.reshape(-1, D)
+    inv = scaler.inverse_transform(flat)
+    return inv.reshape(arr.shape)
+
+
+def compute_scores_scaled_and_orig(trues: np.ndarray, preds: np.ndarray, scaler):
+    score_scaled = evaluate(trues, preds)
+
+    trues_inv = _inverse_if_possible(trues, scaler)
+    preds_inv = _inverse_if_possible(preds, scaler)
+    score_orig = evaluate(trues_inv, preds_inv)
+
+    return score_scaled, score_orig
+
+
+def _fmt_score(tag, score):
+    mape, mae, rmse, r2, rse = score
+    mape_pct = mape * 100.0
+    return f"{tag}: MAPE {mape_pct:10.6f}; MAE {mae:10.6f}; RMSE {rmse:10.6f}; R2 {r2:10.6f}; RSE {rse:10.6f}."
+
+
+# args
+parser = argparse.ArgumentParser(description='SpikF-GO: Spiking Fourier Graph Operators for Multivariate Time Series Forecasting')
+parser.add_argument('--data', type=str, default='ECG', help='data set')
+parser.add_argument('--feature_size', type=int, default=140, help='feature size')
+parser.add_argument('--seq_length', type=int, default=12, help='input length')
+parser.add_argument('--pre_length', type=int, default=12, help='predict length')
+parser.add_argument('--embed_size', type=int, default=128, help='embedding dimensions')
+parser.add_argument('--hidden_size', type=int, default=256, help='hidden dimensions')
+parser.add_argument('--train_epochs', type=int, default=100, help='train epochs')
+parser.add_argument('--batch_size', type=int, default=4, help='input data batch size')
+parser.add_argument('--learning_rate', type=float, default=0.00001, help='optimizer learning rate')
+parser.add_argument('--exponential_decay_step', type=int, default=5)
+parser.add_argument('--validate_freq', type=int, default=1)
+parser.add_argument('--early_stop', type=bool, default=False)
+parser.add_argument('--decay_rate', type=float, default=0.5)
+parser.add_argument('--train_ratio', type=float, default=0.6)
+parser.add_argument('--val_ratio', type=float, default=0.2)
+parser.add_argument('--device', type=str, default='cuda:0', help='device')
+parser.add_argument('--tau', type=float, default=2.0, help='tau')
+parser.add_argument('--alpha', type=float, default=1.0)
+parser.add_argument('--T', type=int, default=4)
+parser.add_argument('--proj_dim', type=int, default=32, help='proj dim')
+parser.add_argument('--model', type=str, default='FGN', help='model name')
+
+parser.add_argument('--patch_num', type=int, default=4)
+parser.add_argument('--patch_dim', type=int, default=16)
+parser.add_argument('--blocks', type=int, default=1)
+parser.add_argument('--energy_loss', type=bool, default=False)
+parser.add_argument('--normalize', action='store_false', help='Disable normalization')
+parser.add_argument('--affine', action='store_false', help='Disable affine layer')
+parser.add_argument('--kernel_size', type=int, default=16)
+
+args = parser.parse_args()
+print(f'Training configs: {args}')
+
+
+data_parser = {
+    'traffic':      {'root_path': 'data/traffic.npy',    'type': '0'},
+    'ECG':          {'root_path': 'data/ECG_data.csv',   'type': '0'},
+    'COVID':        {'root_path': 'data/covid.csv',      'type': '0'},
+    'electricity':  {'root_path': 'data/electricity.csv','type': '0'},
+    'solar':        {'root_path': './data/solar',        'type': '0'},
+    'metr':         {'root_path': 'data/metr.csv',       'type': '0'},
+    'wiki':         {'root_path': 'data/wiki.csv',       'type': '0'},
+    'pems_bay':     {'root_path': 'data/pems-bay.h5',    'type': '0'},
+}
+
+data_dict = {
+    'ECG':         Dataset_ECG,
+    'COVID':       Dataset_ECG,
+    'traffic':     Dataset_Dhfm,
+    'solar':       Dataset_Solar,
+    'wiki':        Dataset_Wiki,
+    'electricity': Dataset_ECG,
+    'metr':        Dataset_ECG,
+    'pems_bay':    Dataset_PEMS_BAY,
+}
+
+if args.data not in data_parser:
+    raise ValueError(f"Unknown dataset {args.data}. Available: {list(data_parser.keys())}")
+
+data_info = data_parser[args.data]
+Data = data_dict[args.data]
+
+
+train_set = Data(
+    root_path=data_info['root_path'], flag='train',
+    seq_len=args.seq_length, pre_len=args.pre_length,
+    type=data_info['type'], train_ratio=args.train_ratio, val_ratio=args.val_ratio,
+    scaler=None
+)
+train_scaler = getattr(train_set, "scaler", None)
+
+val_set = Data(
+    root_path=data_info['root_path'], flag='val',
+    seq_len=args.seq_length, pre_len=args.pre_length,
+    type=data_info['type'], train_ratio=args.train_ratio, val_ratio=args.val_ratio,
+    scaler=train_scaler
+)
+
+test_set = Data(
+    root_path=data_info['root_path'], flag='test',
+    seq_len=args.seq_length, pre_len=args.pre_length,
+    type=data_info['type'], train_ratio=args.train_ratio, val_ratio=args.val_ratio,
+    scaler=train_scaler
+)
+
+train_dataloader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True,  num_workers=0, drop_last=True)
+val_dataloader   = DataLoader(val_set,   batch_size=args.batch_size, shuffle=False, num_workers=0, drop_last=False)
+test_dataloader  = DataLoader(test_set,  batch_size=args.batch_size, shuffle=False, num_workers=0, drop_last=False)
+
+print("Train samples:", len(train_set))
+print("Val samples:", len(val_set))
+print("Test samples:", len(test_set))
+
+
+MODELS_SET2 = ["TSGRU", "TSTCN", "TSFormer", "Spikformer_CPG", "SpikeGRU", "SpikeRNN_CPG", "SpikeTCN_CPG"]
+
+
+
+def validate(model, vali_loader, scaler):
+    model.eval()
+    cnt = 0
+    loss_total = 0.0
+    preds_list = []
+    trues_list = []
+
+    for x, y in vali_loader:
+        if args.model in MODELS_SET2 and args.model != 'TSGRU':
+            reset_states(model=model)
+        elif args.model == 'TSGRU':
+            remove(model=model.net[0].tslif)
+
+        x = x.float().to(args.device)
+        y = y.float().to(args.device)
+
+        forecast, _ = model(x)
+        if len(forecast.shape) == 4:
+            forecast = forecast.mean(dim=0)
+
+        loss = forecast_loss(forecast, y)
+        loss_total += float(loss)
+        cnt += 1
+
+        if args.model not in MODELS_SET2:
+            functional.reset_net(model)
+
+        preds_list.append(forecast.detach().cpu().numpy())
+        trues_list.append(y.detach().cpu().numpy())
+
+    preds = np.concatenate(preds_list, axis=0)
+    trues = np.concatenate(trues_list, axis=0)
+
+    score_scaled, score_orig = compute_scores_scaled_and_orig(trues, preds, scaler)
+
+    print(_fmt_score("SCALED", score_scaled))
+    print(_fmt_score("ORIG  ", score_orig))
+
+    model.train()
+    return loss_total / max(1, cnt)
+
+
+def test(model, result_test_file, scaler, load_epoch=97):
+    model = load_model_ts(model, result_test_file, load_epoch)
+    model.eval()
+
+    preds_list = []
+    trues_list = []
+
+    for x, y in test_dataloader:
+        if args.model in MODELS_SET2 and args.model != 'TSGRU':
+            reset_states(model=model)
+        elif args.model == 'TSGRU':
+            remove(model=model.net[0].tslif)
+
+        x = x.float().to(args.device)
+        y = y.float().to(args.device)
+
+        forecast, _ = model(x)
+        if len(forecast.shape) == 4:
+            forecast = forecast.mean(dim=0)
+
+        if args.model not in MODELS_SET2:
+            functional.reset_net(model)
+
+        preds_list.append(forecast.detach().cpu().numpy())
+        trues_list.append(y.detach().cpu().numpy())
+
+    preds = np.concatenate(preds_list, axis=0)
+    trues = np.concatenate(trues_list, axis=0)
+
+    score_scaled, score_orig = compute_scores_scaled_and_orig(trues, preds, scaler)
+
+    print(_fmt_score("SCALED", score_scaled))
+    print(_fmt_score("ORIG  ", score_orig))
+
+    return score_scaled, score_orig
+
+
+def build_opt_sched(model, lr=3e-4, wd=0.01, gate_lr_ratio=0.3,
+                    warmup_epochs=8, total_epochs=100):
+    decay, no_decay, gate = [], [], []
+    for name, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        name_l = name.lower()
+        is_bias = name.endswith('bias')
+        is_norm = ('norm' in name_l) or ('bn' in name_l)
+        is_embed = ('embeddings' in name_l) or ('time_basis' in name_l)
+        if 'freq_gate' in name_l and 'log_alpha' in name_l:
+            no_decay.append(p)
+        elif is_bias or is_norm or is_embed or p.ndim == 1:
+            no_decay.append(p)
+        else:
+            decay.append(p)
+
+    optim = torch.optim.AdamW([
+        {'params': decay,    'lr': lr, 'weight_decay': wd},
+        {'params': no_decay, 'lr': lr, 'weight_decay': 0.0},
+    ], betas=(0.9, 0.99), eps=1e-8)
+
+    warmup = torch.optim.lr_scheduler.LinearLR(
+        optim, start_factor=0.1, end_factor=1.0, total_iters=warmup_epochs
+    )
+    cosine = torch.optim.lr_scheduler.CosineAnnealingLR(
+        optim, T_max=max(1, total_epochs - warmup_epochs), eta_min=lr * 0.1
+    )
+    sched = torch.optim.lr_scheduler.SequentialLR(
+        optim, schedulers=[warmup, cosine], milestones=[warmup_epochs]
+    )
+    return optim, sched
+
+
+
+if __name__ == '__main__':
+
+    seeds = [2021, 2022, 2023, 2024, 2025]
+
+    scaled_results = {'mape': [], 'mae': [], 'rmse': [], 'r2': [], 'rse': []}
+    orig_results   = {'mape': [], 'mae': [], 'rmse': [], 'r2': [], 'rse': []}
+
+    for run_idx, seed in enumerate(seeds):
+        print(f"\n{'='*60}")
+        print(f"Starting Run {run_idx + 1}/5 | seed={seed}")
+        print(f"{'='*60}")
+
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
+
+        result_train_file = os.path.join('output', args.data, args.model, f'train_run_{run_idx+1}_seed_{seed}')
+        result_test_file  = os.path.join('output', args.data, args.model, f'train_run_{run_idx+1}_seed_{seed}')
+        os.makedirs(result_train_file, exist_ok=True)
+        os.makedirs(result_test_file,  exist_ok=True)
+
+        device = torch.device(args.device if torch.cuda.is_available() else "cpu")
+
+        if args.model == 'SpikF_GO':
+            model = SpikF_GO(args, pre_length=args.pre_length, embed_size=args.embed_size,
+                              feature_size=args.feature_size, seq_length=args.seq_length, hidden_size=args.hidden_size)
+            my_optim, my_lr_scheduler = build_opt_sched(
+                model, lr=args.learning_rate, wd=0.01,
+                warmup_epochs=max(4, args.train_epochs//8), total_epochs=args.train_epochs
+            )
+        elif args.model == 'SpikF_GO_CPG':
+            model = SpikF_GO_CPG(args, pre_length=args.pre_length, embed_size=args.embed_size,
+                                  feature_size=args.feature_size, seq_length=args.seq_length, hidden_size=args.hidden_size)
+            my_optim, my_lr_scheduler = build_opt_sched(
+                model, lr=args.learning_rate, wd=0.01,
+                warmup_epochs=max(4, args.train_epochs//8), total_epochs=args.train_epochs
+            )
+        elif args.model == 'FGN':
+            model = FGN(args, pre_length=args.pre_length, embed_size=args.embed_size,
+                        feature_size=args.feature_size, seq_length=args.seq_length, hidden_size=args.hidden_size)
+            my_optim = torch.optim.RMSprop(params=model.parameters(), lr=args.learning_rate, eps=1e-08)
+            my_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=my_optim, gamma=args.decay_rate)
+        elif args.model == 'SpikF':
+            model = SpikF(args, input_len=args.seq_length, patch_num=args.patch_num, patch_dim=args.patch_dim,
+                          T=args.T, blocks=args.blocks, D=args.feature_size, pred_len=args.pre_length,
+                          tau=args.tau, alpha=args.alpha, hidden_dim=args.hidden_size)
+            my_optim = torch.optim.RMSprop(params=model.parameters(), lr=args.learning_rate, eps=1e-08)
+            my_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=my_optim, gamma=args.decay_rate)
+        elif args.model == 'iSpikformer':
+            model = iSpikformer(args, input_len=args.seq_length, patch_num=args.patch_num, patch_dim=args.patch_dim,
+                                T=args.T, blocks=args.blocks, D=args.feature_size, pred_len=args.pre_length,
+                                tau=args.tau, alpha=args.alpha, hidden_dim=args.hidden_size)
+            my_optim = torch.optim.RMSprop(params=model.parameters(), lr=args.learning_rate, eps=1e-08)
+            my_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=my_optim, gamma=args.decay_rate)
+        elif args.model == 'TSGRU':
+            model = TSGRU(args, hidden_size=args.hidden_size, layers=args.blocks,
+                         num_steps=args.T, input_size=args.feature_size)
+            my_optim = torch.optim.RMSprop(params=model.parameters(), lr=args.learning_rate, eps=1e-08)
+            my_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=my_optim, gamma=args.decay_rate)
+        elif args.model == 'TSTCN':
+            model = TSTCN(args=args, num_levels=args.blocks)
+            my_optim = torch.optim.RMSprop(params=model.parameters(), lr=args.learning_rate, eps=1e-08)
+            my_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=my_optim, gamma=args.decay_rate)
+        elif args.model == 'TSFormer':
+            model = TSFormer(args=args)
+            my_optim = torch.optim.RMSprop(params=model.parameters(), lr=args.learning_rate, eps=1e-08)
+            my_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=my_optim, gamma=args.decay_rate)
+        elif args.model == 'Spikformer_CPG':
+            model = Spikformer_CPG(args=args)
+            my_optim = torch.optim.RMSprop(params=model.parameters(), lr=args.learning_rate, eps=1e-08)
+            my_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=my_optim, gamma=args.decay_rate)
+        elif args.model == 'SpikeGRU':
+            model = SpikeGRU(args, hidden_size=args.hidden_size, layers=args.blocks,
+                             num_steps=args.T, input_size=args.feature_size)
+            my_optim = torch.optim.RMSprop(params=model.parameters(), lr=args.learning_rate, eps=1e-08)
+            my_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=my_optim, gamma=args.decay_rate)
+        elif args.model == 'SpikeRNN_CPG':
+            model = SpikeRNN_CPG(args, hidden_size=args.hidden_size, layers=args.blocks,
+                             num_steps=args.T, input_size=args.feature_size)
+            my_optim = torch.optim.RMSprop(params=model.parameters(), lr=args.learning_rate, eps=1e-08)
+            my_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=my_optim, gamma=args.decay_rate)
+        elif args.model == 'SpikeTCN_CPG':
+            model = SpikeTCN_CPG(args=args, num_levels=args.blocks)
+            my_optim = torch.optim.RMSprop(params=model.parameters(), lr=args.learning_rate, eps=1e-08)
+            my_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=my_optim, gamma=args.decay_rate)
+        else:
+            raise ValueError(f"Unknown model: {args.model}")
+
+        model = model.to(device)
+        forecast_loss = nn.MSELoss(reduction='mean').to(device)
+
+        # train
+        for epoch in range(args.train_epochs):
+            warm = int(0.3 * args.train_epochs)
+            cool = epoch >= warm
+
+            epoch_start_time = time.time()
+            model.train()
+            loss_total = 0.0
+            cnt = 0
+
+            for x, y in train_dataloader:
+                if args.model in MODELS_SET2 and args.model != 'TSGRU':
+                    reset_states(model=model)
+                elif args.model == 'TSGRU':
+                    remove(model=model.net[0].tslif)
+
+                x = x.float().to(device)
+                y = y.float().to(device)
+
+                forecast, aux = model(x)
+
+                if len(forecast.shape) == 4:
+                    y_rep = y.repeat(args.T, 1, 1, 1)
+                else:
+                    y_rep = y
+
+                if (args.model in ['SpikF_GO', 'SpikF_GO_CPG']) and args.energy_loss:
+                    energy_lambda = 20.0
+                    mse = forecast_loss(forecast, y_rep)
+                    adaptive_lambda = (mse.detach() / 100.0) * energy_lambda
+                    loss = mse + adaptive_lambda * aux["rho_hat"]
+                else:
+                    loss = forecast_loss(forecast, y_rep)
+
+                my_optim.zero_grad(set_to_none=True)
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                my_optim.step()
+
+                loss_total += float(loss)
+                cnt += 1
+
+                if args.model not in MODELS_SET2:
+                    functional.reset_net(model)
+
+            if (epoch + 1) % args.exponential_decay_step == 0:
+                my_lr_scheduler.step()
+
+            if (epoch + 1) % args.validate_freq == 0:
+                val_loss = validate(model, val_dataloader, train_scaler)
+                enc_rate_v = float(aux.get('enc_rate', torch.tensor(0.0)))
+                gate_l0_v = float(aux.get('rho_hat', torch.tensor(0.0)))
+                freq_act_v = float(aux.get('freq_mask_active', torch.tensor(0.0)))
+
+                print('Run {} | epoch {:03d} | {:5.2f}s | train_loss {:5.4f} | val_loss {:5.4f} | enc_rate {:.3f} | gate_L0 {:.3f} | f_active {:.3f}'.format(
+                    run_idx + 1, epoch, (time.time() - epoch_start_time), loss_total / max(1, cnt), val_loss,
+                    enc_rate_v, gate_l0_v, freq_act_v))
+
+            save_model_ts(model, result_train_file, epoch)
+
+        save_model_ts(model, result_train_file, f'final_run_{run_idx+1}')
+
+        print("--- TEST ---")
+        score_scaled, score_orig = test(model, result_test_file, train_scaler, load_epoch=97)
+
+        scaled_results['mape'].append(score_scaled[0])
+        scaled_results['mae'].append(score_scaled[1])
+        scaled_results['rmse'].append(score_scaled[2])
+        scaled_results['r2'].append(score_scaled[3])
+        scaled_results['rse'].append(score_scaled[4])
+
+        orig_results['mape'].append(score_orig[0])
+        orig_results['mae'].append(score_orig[1])
+        orig_results['rmse'].append(score_orig[2])
+        orig_results['r2'].append(score_orig[3])
+        orig_results['rse'].append(score_orig[4])
+
+        print(f"Run {run_idx + 1} completed.")
+        print(_fmt_score("Results", score_scaled))
+
+    def _mean_std(arr):
+        arr = np.asarray(arr, dtype=np.float64)
+        return float(np.mean(arr)), float(np.std(arr))
+
+    print(f"\n{'='*60}")
+    print("FINAL RESULTS ACROSS RUNS ")
+    print(f"{'='*60}")
+
+    for tag, store in [("SCALED", scaled_results)]:
+        mape_pct = np.asarray(store['mape'], dtype=np.float64) * 100.0
+        m_mean, m_std = _mean_std(mape_pct)
+        a_mean, a_std = _mean_std(store['mae'])
+        r_mean, r_std = _mean_std(store['rmse'])
+        r2_mean, r2_std = _mean_std(store['r2'])
+        rse_mean, rse_std = _mean_std(store['rse'])
+
+        print(f"\n[{tag}]")
+        print(f"MAPE: {mape_pct}  | mean={m_mean:.6f} std={m_std:.6f}")
+        print(f"MAE : {np.array(store['mae'])}   | mean={a_mean:.6f} std={a_std:.6f}")
+        print(f"RMSE: {np.array(store['rmse'])}  | mean={r_mean:.6f} std={r_std:.6f}")
+        print(f"R2  : {np.array(store['r2'])}    | mean={r2_mean:.6f} std={r2_std:.6f}")
+        print(f"RSE  : {np.array(store['rse'])}    | mean={rse_mean:.6f} std={rse_std:.6f}")
+
+    summary_file = os.path.join('output', args.data, args.model, 'summary_results.txt')
+    os.makedirs(os.path.dirname(summary_file), exist_ok=True)
+
+    with open(summary_file, 'w') as f:
+        f.write("Results across 5 runs:\n")
+        f.write(f"Seeds used: {seeds}\n\n")
+
+        for tag, store in [("SCALED", scaled_results)]:
+            mape_pct = np.asarray(store['mape'], dtype=np.float64) * 100.0
+            m_mean, m_std = _mean_std(mape_pct)
+            a_mean, a_std = _mean_std(store['mae'])
+            r_mean, r_std = _mean_std(store['rmse'])
+            r2_mean, r2_std = _mean_std(store['r2'])
+            rse_mean, rse_std = _mean_std(store['rse'])
+
+            f.write(f"[{tag}]\n")
+            f.write(f"MAPE - Individual: {mape_pct}\n")
+            f.write(f"MAPE - Mean: {m_mean:.6f}, Std: {m_std:.6f}\n")
+            f.write(f"MAE  - Individual: {np.array(store['mae'])}\n")
+            f.write(f"MAE  - Mean: {a_mean:.6f}, Std: {a_std:.6f}\n")
+            f.write(f"RMSE - Individual: {np.array(store['rmse'])}\n")
+            f.write(f"RMSE - Mean: {r_mean:.6f}, Std: {r_std:.6f}\n")
+            f.write(f"R2   - Individual: {np.array(store['r2'])}\n")
+            f.write(f"R2   - Mean: {r2_mean:.6f}, Std: {r2_std:.6f}\n\n")
+            f.write(f"RSE   - Individual: {np.array(store['rse'])}\n")
+            f.write(f"RSE   - Mean: {rse_mean:.6f}, Std: {rse_std:.6f}\n\n")
+
+    print(f"\nSaved summary to: {summary_file}")

utils/utils.py

@@ -0,0 +1,252 @@
+# -*- coding:utf-8 -*-
+
+import numpy as np
+import torch
+import os
+
+
+def concat_fun(inputs, axis=-1):
+    if len(inputs) == 1:
+        return inputs[0]
+    else:
+        return torch.cat(inputs, dim=axis)
+
+
+def slice_arrays(arrays, start=None, stop=None):
+    """Slice an array or list of arrays.
+
+    This takes an array-like, or a list of
+    array-likes, and outputs:
+        - arrays[start:stop] if `arrays` is an array-like
+        - [x[start:stop] for x in arrays] if `arrays` is a list
+
+    Can also work on list/array of indices: `slice_arrays(x, indices)`
+
+    Arguments:
+        arrays: Single array or list of arrays.
+        start: can be an integer index (start index)
+            or a list/array of indices
+        stop: integer (stop index); should be None if
+            `start` was a list.
+
+    Returns:
+        A slice of the array(s).
+
+    Raises:
+        ValueError: If the value of start is a list and stop is not None.
+    """
+
+    if arrays is None:
+        return [None]
+
+    if isinstance(arrays, np.ndarray):
+        arrays = [arrays]
+
+    if isinstance(start, list) and stop is not None:
+        raise ValueError('The stop argument has to be None if the value of start '
+                         'is a list.')
+    elif isinstance(arrays, list):
+        if hasattr(start, '__len__'):
+            # hdf5 datasets only support list objects as indices
+            if hasattr(start, 'shape'):
+                start = start.tolist()
+            return [None if x is None else x[start] for x in arrays]
+        else:
+            if len(arrays) == 1:
+                return arrays[0][start:stop]
+            return [None if x is None else x[start:stop] for x in arrays]
+    else:
+        if hasattr(start, '__len__'):
+            if hasattr(start, 'shape'):
+                start = start.tolist()
+            return arrays[start]
+        elif hasattr(start, '__getitem__'):
+            return arrays[start:stop]
+        else:
+            return [None]
+
+
+def save_model(model, model_dir, epoch=None):
+    if model_dir is None:
+        return
+    if not os.path.exists(model_dir):
+        os.makedirs(model_dir)
+    epoch = str(epoch) if epoch else ''
+    file_name = os.path.join(model_dir, epoch + '_dhfm.pt')
+    with open(file_name, 'wb') as f:
+        torch.save(model, f)
+
+
+def load_model(model_dir, epoch=None):
+    if not model_dir:
+        return
+    epoch = str(epoch) if epoch else ''
+    file_name = os.path.join(model_dir, epoch + '_dhfm.pt')
+    if not os.path.exists(model_dir):
+        os.makedirs(model_dir)
+    if not os.path.exists(file_name):
+        return
+    with open(file_name, 'rb') as f:
+        model = torch.load(f)
+    return model
+
+def masked_MAPE(v, v_, axis=None):
+    '''
+    Mean absolute percentage error.
+    :param v: np.ndarray or int, ground truth.
+    :param v_: np.ndarray or int, prediction.
+    :param axis: axis to do calculation.
+    :return: int, MAPE averages on all elements of input.
+    '''
+    mask = (v == 0)
+    percentage = np.abs(v_ - v) / np.abs(v)
+    if np.any(mask):
+        masked_array = np.ma.masked_array(percentage, mask=mask)  # mask the dividing-zero as invalid
+        result = masked_array.mean(axis=axis)
+        if isinstance(result, np.ma.MaskedArray):
+            return result.filled(np.nan)
+        else:
+            return result
+    return np.mean(percentage, axis).astype(np.float64)
+
+"""
+original
+def MAPE(v, v_, axis=None):
+    '''
+    Mean absolute percentage error.
+    :param v: np.ndarray or int, ground truth.
+    :param v_: np.ndarray or int, prediction.
+    :param axis: axis to do calculation.
+    :return: int, MAPE averages on all elements of input.
+    '''
+    mape = (np.abs(v_ - v) / np.abs(v)+1e-5).astype(np.float64)
+    mape = np.where(mape > 5, 5, mape)
+    return np.mean(mape, axis)
+
+"""
+
+def MAPE(v, v_, axis=None):
+    '''
+    Mean absolute percentage error.
+    :param v: np.ndarray or int, ground truth.
+    :param v_: np.ndarray or int, prediction.
+    :param axis: axis to do calculation.
+    :return: float, MAPE averages on all elements of input.
+    '''
+    mape = (np.abs(v_ - v) / (np.abs(v) + 1e-5)).astype(np.float64)
+    mape = np.where(mape > 5, 5, mape)  # clip extreme values
+    return np.mean(mape, axis)
+
+
+#def MAPE(true, pred):
+#    return np.mean(np.abs((pred - true) / (true+1e-5)))
+
+def smape(P, A):
+    nz = np.where(A > 0)
+    Pz = P[nz]
+    Az = A[nz]
+
+    return np.mean(2 * np.abs(Az - Pz) / (np.abs(Az) + np.abs(Pz)))
+
+
+def R2(y, y_hat, axis=None, eps=1e-12):
+    """
+    R^2 score for arrays shaped like [count, time_step, node] (or compatible).
+    axis=None -> global scalar R2 over all elements.
+    axis can be int or tuple of ints: reduce over those axes, keeping the others.
+    """
+    y = np.asarray(y, dtype=np.float64)
+    y_hat = np.asarray(y_hat, dtype=np.float64)
+
+    # residual sum of squares
+    ss_res = np.sum((y - y_hat) ** 2, axis=axis)
+
+    # total sum of squares around mean of y along the same reduction axis
+    y_mean = np.mean(y, axis=axis, keepdims=True)
+    ss_tot = np.sum((y - y_mean) ** 2, axis=axis)
+
+    # Avoid division by zero (constant targets)
+    denom = ss_tot + eps
+    r2 = 1.0 - (ss_res / denom)
+
+    # If ss_tot is truly ~0, R2 is not well-defined; mark as nan
+    # (Optional) If you want 0.0 instead, replace np.nan with 0.0
+    if np.isscalar(ss_tot):
+        if ss_tot < eps:
+            return np.nan
+        return float(r2)
+
+    r2 = np.where(ss_tot < eps, np.nan, r2)
+    return r2.astype(np.float64)
+
+def RSE(v, v_, axis=None, eps=1e-12):
+    '''
+    Relative squared error (rooted):
+        sqrt( sum((v_ - v)^2) / sum((v - mean(v))^2) )
+    :param v: np.ndarray or int, ground truth.
+    :param v_: np.ndarray or int, prediction.
+    :param axis: axis to do calculation.
+    :return: float, RSE on all elements of input (or reduced by axis).
+    '''
+    v = np.asarray(v, dtype=np.float64)
+    v_ = np.asarray(v_, dtype=np.float64)
+
+    v_mean = np.mean(v, axis=axis, keepdims=True)
+    num = np.sum((v_ - v) ** 2, axis=axis)
+    denom = np.sum((v - v_mean) ** 2, axis=axis)
+    return np.sqrt(num / (denom + eps)).astype(np.float64)
+
+def RMSE(v, v_, axis=None):
+    '''
+    Mean squared error.
+    :param v: np.ndarray or int, ground truth.
+    :param v_: np.ndarray or int, prediction.
+    :param axis: axis to do calculation.
+    :return: int, RMSE averages on all elements of input.
+    '''
+    return np.sqrt(np.mean((v_ - v) ** 2, axis)).astype(np.float64)
+
+
+def MAE(v, v_, axis=None):
+    '''
+    Mean absolute error.
+    :param v: np.ndarray or int, ground truth.
+    :param v_: np.ndarray or int, prediction.
+    :param axis: axis to do calculation.
+    :return: int, MAE averages on all elements of input.
+    '''
+    return np.mean(np.abs(v_ - v), axis).astype(np.float64)
+
+
+def evaluate(y, y_hat, by_step=False, by_node=False):
+    '''
+    :param y: array in shape of [count, time_step, node].
+    :param y_hat: in same shape with y.
+    :param by_step: evaluate by time_step dim.
+    :param by_node: evaluate by node dim.
+    :return: array of mape, mae and rmse.
+    '''
+    if not by_step and not by_node:
+        return MAPE(y, y_hat), MAE(y, y_hat), RMSE(y, y_hat), R2(y, y_hat), RSE(y, y_hat)
+    if by_step and by_node:
+        return MAPE(y, y_hat, axis=0), MAE(y, y_hat, axis=0), RMSE(y, y_hat, axis=0), R2(y, y_hat, axis=0)
+    if by_step:
+        return MAPE(y, y_hat, axis=(0, 2)), MAE(y, y_hat, axis=(0, 2)), RMSE(y, y_hat, axis=(0, 2)), R2(y, y_hat, axis=(0, 2))
+    if by_node:
+        return MAPE(y, y_hat, axis=(0, 1)), MAE(y, y_hat, axis=(0, 1)), RMSE(y, y_hat, axis=(0, 1)), R2(y, y_hat, axis=(0, 1))
+
+
+def save_model_ts(model, path, epoch):
+    if not os.path.exists(path):
+        os.makedirs(path)
+    filename = 'epoch_{}.pth'.format(epoch)
+    f = os.path.join(path, filename)
+    # Save state_dict instead of the entire model
+    torch.save(model.state_dict(), f)
+
+def load_model_ts(model, path, epoch):
+    """Load state dict into an existing model instance"""
+    filename = 'epoch_{}.pth'.format(epoch)
+    f = os.path.join(path, filename)
+    model.load_state_dict(torch.load(f))
+    return model