1{
2  "op": "replace",
3  "filename": "Time-Series-Library/models/Custom.py",
4  "start_line": 1,
5  "end_line": 41,
6  "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom layers.Embed import DataEmbedding\nfrom layers.Transformer_EncDec import Encoder, EncoderLayer\nfrom layers.SelfAttention_Family import FullAttention, AttentionLayer\n\n\nclass Model(nn.Module):\n    \"\"\"\n    Mask-Conditioned Spectral Transformer (MCST) for Time Series Imputation.\n\n    Novel algorithmic contributions:\n\n    1. Mask Conditioning: The binary observation mask [B, L, C] is linearly projected\n       into the d_model embedding space and added to the input representation. This\n       explicitly informs the transformer about which positions are observed vs. masked,\n       enabling it to differentiate between true zeros and masked-to-zero values.\n\n    2. Spectral Residual Correction: After the transformer encoder, a learnable\n       complex-valued frequency filter is applied in the Fourier domain. This adds a\n       residual correction that captures global periodic structures which the local\n       attention mechanism may underweight. Initialized at zero (no-op), it learns\n       useful spectral corrections during training.\n\n    Forward signature: forward(x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None)\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.seq_len  # For imputation, pred_len = seq_len\n        self.enc_in = configs.enc_in\n        d_model = configs.d_model\n\n        # Standard data embedding (value conv + position + temporal features)\n        self.enc_embedding = DataEmbedding(\n            configs.enc_in, d_model, configs.embed, configs.freq, configs.dropout)\n\n        # Novel 1: Mask conditioning — project binary mask to embedding space\n        self.mask_proj = nn.Linear(configs.enc_in, d_model, bias=False)\n\n        # Transformer encoder with bidirectional (non-causal) attention\n        self.encoder = Encoder(\n            [EncoderLayer(\n                AttentionLayer(\n                    FullAttention(False, configs.factor,\n                                 attention_dropout=configs.dropout,\n                                 output_attention=False),\n                    d_model, configs.n_heads),\n                d_model, configs.d_ff,\n                dropout=configs.dropout,\n                activation=configs.activation\n            ) for _ in range(configs.e_layers)],\n            norm_layer=nn.LayerNorm(d_model)\n        )\n\n        # Novel 2: Spectral residual correction (complex-valued frequency filter)\n        n_freq = configs.seq_len // 2 + 1\n        self.spec_real = nn.Parameter(torch.zeros(1, n_freq, d_model))\n        self.spec_imag = nn.Parameter(torch.zeros(1, n_freq, d_model))\n\n        # Output projection\n        c_out = getattr(configs, 'c_out', configs.enc_in)\n        self.projection = nn.Linear(d_model, c_out, bias=True)\n\n    def spectral_refine(self, x):\n        \"\"\"Learnable frequency-domain filter producing a residual correction.\"\"\"\n        x_ft = torch.fft.rfft(x, dim=1)  # [B, n_freq, d_model], complex\n        weight = torch.complex(self.spec_real, self.spec_imag)\n        correction_ft = x_ft * weight\n        return torch.fft.irfft(correction_ft, n=x.shape[1], dim=1)\n\n    def imputation(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask):\n        \"\"\"\n        Imputation: reconstruct missing values in the input sequence.\n        Input: x_enc [batch, seq_len, enc_in] with zeros at masked positions\n        Mask: [batch, seq_len, enc_in], 1=observed, 0=masked\n        Output: [batch, seq_len, enc_in]\n        \"\"\"\n        # --- Mask-Aware Normalization (standard for imputation) ---\n        means = torch.sum(x_enc, dim=1) / torch.sum(mask == 1, dim=1)\n        means = means.unsqueeze(1).detach()\n        x_enc = x_enc - means\n        x_enc = x_enc.masked_fill(mask == 0, 0)\n        stdev = torch.sqrt(torch.sum(x_enc * x_enc, dim=1) /\n                           torch.sum(mask == 1, dim=1) + 1e-5)\n        stdev = stdev.unsqueeze(1).detach()\n        x_enc = x_enc / stdev\n\n        # --- Embedding with Mask Conditioning ---\n        enc_out = self.enc_embedding(x_enc, x_mark_enc)  # [B, L, d_model]\n        enc_out = enc_out + self.mask_proj(mask.float())   # Add mask info\n\n        # --- Transformer Encoder ---\n        enc_out, _ = self.encoder(enc_out)\n\n        # --- Spectral Residual Correction ---\n        enc_out = enc_out + self.spectral_refine(enc_out)\n\n        # --- Output Projection ---\n        dec_out = self.projection(enc_out)\n\n        # --- De-Normalization ---\n        dec_out = dec_out * stdev + means\n        return dec_out\n\n    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None):\n        if self.task_name == 'imputation':\n            dec_out = self.imputation(x_enc, x_mark_enc, x_dec, x_mark_dec, mask)\n            return dec_out\n        return None\n"
7}

1{
2  "op": "replace",
3  "filename": "Time-Series-Library/models/Custom.py",
4  "start_line": 1,
5  "end_line": 113,
6  "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom layers.Embed import DataEmbedding\n\n\nclass ImputationAttention(nn.Module):\n    \"\"\"\n    Novel attention mechanism for imputation with two learnable structural biases:\n    \n    1. Distance bias: Per-head log-linear decay preferring nearby positions.\n       For imputation, nearby observed values are most informative for filling gaps.\n       Each head learns its own decay rate, allowing specialization (some heads \n       focus locally, others globally).\n    \n    2. Mask bias: Per-head additive bias toward observed (reliable) positions.\n       This directs information flow from observed values to masked positions,\n       avoiding the propagation of unreliable (zeroed-out) information.\n    \"\"\"\n    def __init__(self, d_model, n_heads, seq_len, dropout=0.1):\n        super().__init__()\n        self.n_heads = n_heads\n        self.d_k = d_model // n_heads\n        self.scale = self.d_k ** -0.5\n        self.q_proj = nn.Linear(d_model, d_model)\n        self.k_proj = nn.Linear(d_model, d_model)\n        self.v_proj = nn.Linear(d_model, d_model)\n        self.out_proj = nn.Linear(d_model, d_model)\n        self.attn_drop = nn.Dropout(dropout)\n\n        # Per-head learnable bias magnitudes\n        self.mask_bias = nn.Parameter(torch.ones(1, n_heads, 1, 1) * 0.5)\n        self.dist_scale = nn.Parameter(torch.ones(1, n_heads, 1, 1) * 0.05)\n\n        # Precompute log-distance matrix (registered as buffer, not parameter)\n        pos = torch.arange(seq_len).float()\n        dist = (pos.unsqueeze(0) - pos.unsqueeze(1)).abs()\n        self.register_buffer('log_dist', torch.log(dist + 1.0))\n\n    def forward(self, x, mask_ratio=None):\n        B, L, D = x.shape\n        H = self.n_heads\n        Q = self.q_proj(x).view(B, L, H, self.d_k).transpose(1, 2)\n        K = self.k_proj(x).view(B, L, H, self.d_k).transpose(1, 2)\n        V = self.v_proj(x).view(B, L, H, self.d_k).transpose(1, 2)\n\n        # Standard scaled dot-product scores\n        scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale\n\n        # Distance bias: prefer nearby positions (log-linear decay)\n        scores = scores - self.dist_scale * self.log_dist[:L, :L]\n\n        # Mask bias: prefer observed positions\n        if mask_ratio is not None:\n            scores = scores + self.mask_bias * mask_ratio.unsqueeze(1).unsqueeze(2)\n\n        attn = self.attn_drop(torch.softmax(scores, dim=-1))\n        out = torch.matmul(attn, V).transpose(1, 2).contiguous().view(B, L, D)\n        return self.out_proj(out)\n\n\nclass EncoderBlock(nn.Module):\n    \"\"\"Transformer encoder block using ImputationAttention.\"\"\"\n    def __init__(self, d_model, n_heads, d_ff, seq_len, dropout=0.1, activation='gelu'):\n        super().__init__()\n        self.attn = ImputationAttention(d_model, n_heads, seq_len, dropout)\n        self.conv1 = nn.Conv1d(d_model, d_ff, 1)\n        self.conv2 = nn.Conv1d(d_ff, d_model, 1)\n        self.norm1 = nn.LayerNorm(d_model)\n        self.norm2 = nn.LayerNorm(d_model)\n        self.dropout = nn.Dropout(dropout)\n        self.activation = F.gelu if activation == 'gelu' else F.relu\n\n    def forward(self, x, mask_ratio=None):\n        new_x = self.attn(x, mask_ratio)\n        x = x + self.dropout(new_x)\n        y = x = self.norm1(x)\n        y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))\n        y = self.dropout(self.conv2(y).transpose(-1, 1))\n        return self.norm2(x + y)\n\n\nclass Model(nn.Module):\n    \"\"\"\n    Imputation-Specialized Spectral Transformer (ISST).\n\n    Novel algorithmic contributions:\n\n    1. Imputation Attention with Dual Structural Biases:\n       - Distance bias: log-linear decay preferring nearby positions (per head)\n       - Mask bias: additive preference for observed positions (per head)\n       These biases encode the key inductive priors for imputation: locality\n       and reliability of observations.\n\n    2. Mask Conditioning: Binary observation mask projected into embedding space\n       and added to transformer input, encoding the full per-channel missingness\n       pattern at each time step.\n\n    3. Spectral Residual Correction: Learnable complex-valued frequency filter\n       applied after the transformer. Captures global periodic structures as a\n       lightweight residual correction (initialized at zero = no-op).\n\n    4. Eval-Time Iterative Refinement: During inference only, a second forward\n       pass is performed using first-pass predictions to fill masked positions.\n       This propagates information through previously-masked gaps at zero extra\n       training cost.\n\n    Forward signature: forward(x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None)\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.seq_len\n        self.enc_in = configs.enc_in\n        d_model = configs.d_model\n\n        # Data embedding (value conv + position + temporal)\n        self.enc_embedding = DataEmbedding(\n            configs.enc_in, d_model, configs.embed, configs.freq, configs.dropout)\n\n        # Mask conditioning\n        self.mask_proj = nn.Linear(configs.enc_in, d_model, bias=False)\n\n        # Encoder with imputation-specialized attention\n        activation = getattr(configs, 'activation', 'gelu')\n        self.blocks = nn.ModuleList([\n            EncoderBlock(d_model, configs.n_heads, configs.d_ff,\n                         configs.seq_len, configs.dropout, activation)\n            for _ in range(configs.e_layers)\n        ])\n        self.encoder_norm = nn.LayerNorm(d_model)\n\n        # Spectral residual correction\n        n_freq = configs.seq_len // 2 + 1\n        self.spec_real = nn.Parameter(torch.zeros(1, n_freq, d_model))\n        self.spec_imag = nn.Parameter(torch.zeros(1, n_freq, d_model))\n\n        # Output projection\n        c_out = getattr(configs, 'c_out', configs.enc_in)\n        self.projection = nn.Linear(d_model, c_out, bias=True)\n\n    def spectral_refine(self, x):\n        \"\"\"Learnable frequency-domain filter producing a residual correction.\"\"\"\n        x_ft = torch.fft.rfft(x, dim=1)\n        w = torch.complex(self.spec_real, self.spec_imag)\n        return torch.fft.irfft(x_ft * w, n=x.shape[1], dim=1)\n\n    def encode_and_predict(self, current, x_mark_enc, mask_emb, mask_ratio):\n        \"\"\"Single forward pass: embed → encode → spectral refine → project.\"\"\"\n        h = self.enc_embedding(current, x_mark_enc) + mask_emb\n        for blk in self.blocks:\n            h = blk(h, mask_ratio)\n        h = self.encoder_norm(h)\n        h = h + self.spectral_refine(h)\n        return self.projection(h)\n\n    def imputation(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask):\n        \"\"\"\n        Imputation with mask-aware normalization, dual-bias attention,\n        and optional eval-time iterative refinement.\n        \"\"\"\n        # --- Mask-Aware Normalization ---\n        means = torch.sum(x_enc, dim=1) / torch.sum(mask == 1, dim=1)\n        means = means.unsqueeze(1).detach()\n        x_enc = x_enc - means\n        x_enc = x_enc.masked_fill(mask == 0, 0)\n        stdev = torch.sqrt(torch.sum(x_enc * x_enc, dim=1) /\n                           torch.sum(mask == 1, dim=1) + 1e-5)\n        stdev = stdev.unsqueeze(1).detach()\n        x_norm = x_enc / stdev\n\n        # Precompute mask embeddings\n        mask_emb = self.mask_proj(mask.float())\n        mask_ratio = mask.float().mean(dim=-1)  # [B, L] per-timestep obs ratio\n\n        # First pass\n        dec_out = self.encode_and_predict(x_norm, x_mark_enc, mask_emb, mask_ratio)\n\n        # Eval-time iterative refinement (zero training cost)\n        if not self.training:\n            filled = mask * x_norm + (1 - mask) * dec_out\n            dec_out = self.encode_and_predict(filled, x_mark_enc, mask_emb, mask_ratio)\n\n        # De-normalize\n        dec_out = dec_out * stdev + means\n        return dec_out\n\n    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None):\n        if self.task_name == 'imputation':\n            dec_out = self.imputation(x_enc, x_mark_enc, x_dec, x_mark_dec, mask)\n            return dec_out\n        return None\n"
7}

1{
2  "op": "replace",
3  "filename": "Time-Series-Library/models/Custom.py",
4  "start_line": 159,
5  "end_line": 194,
6  "content": "    def imputation(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask):\n        \"\"\"\n        Imputation with mask-aware normalization and dual-bias attention.\n        \"\"\"\n        # --- Mask-Aware Normalization ---\n        means = torch.sum(x_enc, dim=1) / torch.sum(mask == 1, dim=1)\n        means = means.unsqueeze(1).detach()\n        x_enc = x_enc - means\n        x_enc = x_enc.masked_fill(mask == 0, 0)\n        stdev = torch.sqrt(torch.sum(x_enc * x_enc, dim=1) /\n                           torch.sum(mask == 1, dim=1) + 1e-5)\n        stdev = stdev.unsqueeze(1).detach()\n        x_enc = x_enc / stdev\n\n        # Precompute mask embeddings\n        mask_emb = self.mask_proj(mask.float())\n        mask_ratio = mask.float().mean(dim=-1)  # [B, L] per-timestep obs ratio\n\n        # Single pass: embed → encode → spectral refine → project\n        dec_out = self.encode_and_predict(x_enc, x_mark_enc, mask_emb, mask_ratio)\n\n        # De-normalize\n        dec_out = dec_out * stdev + means\n        return dec_out\n\n    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None):\n        if self.task_name == 'imputation':\n            dec_out = self.imputation(x_enc, x_mark_enc, x_dec, x_mark_dec, mask)\n            return dec_out\n        return None\n"
7}