1{
2  "end_line": 41,
3  "op": "replace",
4  "start_line": 1,
5  "filename": "Time-Series-Library/models/Custom.py",
6  "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass SpectralMoEBlock(nn.Module):\n    def __init__(self, seq_len, pred_len, num_experts=3):\n        super(SpectralMoEBlock, self).__init__()\n        self.seq_len = seq_len\n        self.pred_len = pred_len\n        self.num_experts = num_experts\n        freq_len = seq_len // 2 + 1\n        \n        self.gate = nn.Sequential(\n            nn.Linear(freq_len, 64),\n            nn.GELU(),\n            nn.Linear(64, freq_len * num_experts)\n        )\n        \n        self.forecast_experts = nn.ModuleList([\n            nn.Linear(seq_len, pred_len) for _ in range(num_experts)\n        ])\n        \n        self.backcast_experts = nn.ModuleList([\n            nn.Linear(seq_len, seq_len) for _ in range(num_experts)\n        ])\n\n    def forward(self, x):\n        # x: [B, D, L]\n        batch_size, D, L = x.shape\n        freq_len = L // 2 + 1\n        \n        # FFT\n        xfft = torch.fft.rfft(x, dim=-1) # [B, D, freq_len]\n        amp = torch.abs(xfft) # [B, D, freq_len]\n        \n        # Gating\n        logits = self.gate(amp) # [B, D, freq_len * num_experts]\n        logits = logits.view(batch_size, D, freq_len, self.num_experts)\n        gates = torch.softmax(logits, dim=-1) # [B, D, freq_len, num_experts]\n        \n        # Expert specific signals\n        y_forecast = 0\n        x_backcast = 0\n        \n        for k in range(self.num_experts):\n            # Masking in freq domain\n            mask_k = gates[..., k] # [B, D, freq_len]\n            xfft_k = xfft * mask_k\n            \n            # Inverse FFT\n            x_k = torch.fft.irfft(xfft_k, n=self.seq_len, dim=-1) # [B, D, L]\n            \n            # Linear projection\n            y_forecast = y_forecast + self.forecast_experts[k](x_k)\n            x_backcast = x_backcast + self.backcast_experts[k](x_k)\n            \n        return x_backcast, y_forecast\n\nclass Model(nn.Module):\n    \"\"\"\n    Custom model for short-term time series forecasting (M4 dataset).\n    Incorporates Spectral Mixture of Experts (Spectral-MoE) on continuous residual blocks.\n    Algorithm:\n    1. Input is Instance Normalized.\n    2. Deep Residual Blocks sequentially operate on the series.\n    3. Within each block, a dynamic frequency-domain router (MoE) isolates\n       sub-spaces (e.g. dynamic low/high pass filtering) learned from the signal amplitude.\n    4. Time-domain linear experts specialize in predicting future (forecast) \n       and reconstructing the input (backcast).\n    \"\"\"\n\n    def __init__(self, configs):\n        super(Model, self).__init__()\n        self.task_name = configs.task_name\n        self.seq_len = configs.seq_len\n        self.pred_len = configs.pred_len\n        self.enc_in = configs.enc_in\n        self.c_out = configs.c_out\n        \n        self.num_blocks = 3  # Deep residual stages\n        num_experts = 3      # Specialists per stage\n        \n        self.blocks = nn.ModuleList([\n            SpectralMoEBlock(self.seq_len, self.pred_len, num_experts=num_experts)\n            for _ in range(self.num_blocks)\n        ])\n        \n        if self.c_out != self.enc_in:\n            self.projection = nn.Linear(self.enc_in, self.c_out)\n        else:\n            self.projection = None\n\n    def forecast(self, x_enc, x_mark_enc, x_dec, x_mark_dec):\n        \"\"\"\n        Forecasting logic.\n        Input: x_enc [batch, seq_len, enc_in]\n        Output: [batch, pred_len, c_out]\n        \"\"\"\n        # RevIN-like Instance Normalization\n        means = x_enc.mean(dim=1, keepdim=True).detach()\n        x_enc_norm = x_enc - means\n        stdev = torch.sqrt(torch.var(x_enc_norm, dim=1, keepdim=True, unbiased=False) + 1e-5)\n        x_enc_norm = x_enc_norm / stdev\n        \n        # Reshape to [B, D, L]\n        x = x_enc_norm.transpose(1, 2)\n        \n        y_pred = torch.zeros(x.shape[0], x.shape[1], self.pred_len, device=x.device)\n        \n        for block in self.blocks:\n            x_backcast, block_forecast = block(x)\n            x = x - x_backcast\n            y_pred = y_pred + block_forecast\n            \n        # Reshape back to [B, L, D]\n        y_pred = y_pred.transpose(1, 2)\n        \n        # De-Normalization\n        y_pred = y_pred * (stdev[:, 0, :].unsqueeze(1).repeat(1, self.pred_len, 1))\n        y_pred = y_pred + (means[:, 0, :].unsqueeze(1).repeat(1, self.pred_len, 1))\n        \n        if self.projection is not None:\n            y_pred = self.projection(y_pred)\n            \n        return y_pred\n\n    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None):\n        if self.task_name == 'long_term_forecast' or self.task_name == 'short_term_forecast':\n            dec_out = self.forecast(x_enc, x_mark_enc, x_dec, x_mark_dec)\n            return dec_out[:, -self.pred_len:, :]\n        return None\n"
7}

1{
2  "end_line": 81,
3  "start_line": 80,
4  "content": "        self.num_blocks = 2  # Deep residual stages\n        num_experts = 4      # Specialists per stage\n",
5  "filename": "Time-Series-Library/models/Custom.py",
6  "op": "replace"
7}

1{
2  "end_line": 131,
3  "start_line": 1,
4  "filename": "Time-Series-Library/models/Custom.py",
5  "op": "replace",
6  "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom layers.Autoformer_EncDec import series_decomp\n\nclass SpectralMoEBlock(nn.Module):\n    def __init__(self, seq_len, pred_len, num_experts=3):\n        super(SpectralMoEBlock, self).__init__()\n        self.seq_len = seq_len\n        self.pred_len = pred_len\n        self.num_experts = num_experts\n        freq_len = seq_len // 2 + 1\n        \n        self.gate = nn.Sequential(\n            nn.Linear(freq_len, 64),\n            nn.GELU(),\n            nn.Dropout(0.1),\n            nn.Linear(64, freq_len * num_experts)\n        )\n        \n        self.forecast_experts = nn.ModuleList([\n            nn.Linear(seq_len, pred_len) for _ in range(num_experts)\n        ])\n        \n        self.backcast_experts = nn.ModuleList([\n            nn.Linear(seq_len, seq_len) for _ in range(num_experts)\n        ])\n\n    def forward(self, x):\n        batch_size, D, L = x.shape\n        freq_len = L // 2 + 1\n        \n        xfft = torch.fft.rfft(x, dim=-1) \n        amp = torch.abs(xfft) \n        \n        logits = self.gate(amp) \n        logits = logits.view(batch_size, D, freq_len, self.num_experts)\n        gates = torch.softmax(logits, dim=-1) \n        \n        y_forecast = 0\n        x_backcast = 0\n        \n        for k in range(self.num_experts):\n            mask_k = gates[..., k] \n            xfft_k = xfft * mask_k\n            \n            x_k = torch.fft.irfft(xfft_k, n=self.seq_len, dim=-1) \n            \n            y_forecast = y_forecast + self.forecast_experts[k](x_k)\n            x_backcast = x_backcast + self.backcast_experts[k](x_k)\n            \n        return x_backcast, y_forecast\n\nclass Model(nn.Module):\n    def __init__(self, configs):\n        super(Model, self).__init__()\n        self.task_name = configs.task_name\n        self.seq_len = configs.seq_len\n        self.pred_len = configs.pred_len\n        self.enc_in = configs.enc_in\n        self.c_out = configs.c_out\n        \n        self.decomp = series_decomp(25)\n        self.linear_trend = nn.Linear(self.seq_len, self.pred_len)\n        \n        self.num_blocks = 3  \n        num_experts = 3      \n        \n        self.blocks = nn.ModuleList([\n            SpectralMoEBlock(self.seq_len, self.pred_len, num_experts=num_experts)\n            for _ in range(self.num_blocks)\n        ])\n        \n        if self.c_out != self.enc_in:\n            self.projection = nn.Linear(self.enc_in, self.c_out)\n        else:\n            self.projection = None\n\n    def forecast(self, x_enc, x_mark_enc, x_dec, x_mark_dec):\n        # RevIN-like Instance Normalization\n        means = x_enc.mean(dim=1, keepdim=True).detach()\n        x_enc_norm = x_enc - means\n        stdev = torch.sqrt(torch.var(x_enc_norm, dim=1, keepdim=True, unbiased=False) + 1e-5)\n        x_enc_norm = x_enc_norm / stdev\n        \n        # Decomposition\n        res, trend = self.decomp(x_enc_norm) # [B, L, D]\n        \n        y_trend = self.linear_trend(trend.transpose(1, 2)).transpose(1, 2)\n        \n        # Reshape to [B, D, L] for Spectral blocks\n        x = res.transpose(1, 2)\n        \n        y_sea = torch.zeros(x.shape[0], x.shape[1], self.pred_len, device=x.device)\n        \n        for block in self.blocks:\n            x_backcast, block_forecast = block(x)\n            x = x - x_backcast\n            y_sea = y_sea + block_forecast\n            \n        y_sea = y_sea.transpose(1, 2) # [B, pred_len, D]\n        \n        y_pred = y_trend + y_sea\n        \n        # De-Normalization\n        y_pred = y_pred * (stdev[:, 0, :].unsqueeze(1).repeat(1, self.pred_len, 1))\n        y_pred = y_pred + (means[:, 0, :].unsqueeze(1).repeat(1, self.pred_len, 1))\n        \n        if self.projection is not None:\n            y_pred = self.projection(y_pred)\n            \n        return y_pred\n\n    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None):\n        if self.task_name == 'long_term_forecast' or self.task_name == 'short_term_forecast':\n            dec_out = self.forecast(x_enc, x_mark_enc, x_dec, x_mark_dec)\n            return dec_out[:, -self.pred_len:, :]\n        return None"
7}