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\n\n\nclass Model(nn.Module):\n    \"\"\"\n    Custom model for time series anomaly detection.\n    Memory-augmented autoencoder for anomaly detection.\n    \n    Forward signature: forward(x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None)\n    - x_enc: [batch, seq_len, enc_in] — input time series\n    - x_mark_enc: not used for anomaly detection (None)\n    - x_dec: not used for anomaly detection (None)\n    - x_mark_dec: not used for anomaly detection (None)\n\n    Must return: [batch, seq_len, c_out] — reconstructed sequence\n    The framework computes MSE between input and output for anomaly scoring.\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 anomaly detection, pred_len = seq_len\n        self.enc_in = configs.enc_in\n        self.c_out = configs.c_out\n        \n        # Hyperparameters\n        self.d_model = 64\n        self.memory_size = 32\n        self.num_layers = 2\n        self.kernel_size = 3\n        self.dropout = 0.1\n        \n        # Encoder: temporal convolutional layers\n        self.encoder = nn.Sequential(\n            nn.Conv1d(self.enc_in, self.d_model, kernel_size=self.kernel_size, padding='same'),\n            nn.ReLU(),\n            nn.Dropout(self.dropout),\n            nn.Conv1d(self.d_model, self.d_model, kernel_size=self.kernel_size, padding='same'),\n            nn.ReLU(),\n            nn.Dropout(self.dropout),\n        )\n        \n        # Positional encoding\n        self.pos_encoder = PositionalEncoding(self.d_model, max_len=self.seq_len)\n        \n        # Transformer layer for temporal dependencies\n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=self.d_model,\n            nhead=4,\n            dim_feedforward=128,\n            dropout=self.dropout,\n            activation='relu',\n            batch_first=True\n        )\n        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=self.num_layers)\n        \n        # Memory module\n        self.memory = nn.Parameter(torch.randn(self.memory_size, self.d_model))\n        self.attention = nn.MultiheadAttention(self.d_model, num_heads=4, dropout=self.dropout, batch_first=True)\n        \n        # Decoder: reconstruct from memory-enhanced representation\n        self.decoder = nn.Sequential(\n            nn.Conv1d(self.d_model, self.d_model, kernel_size=self.kernel_size, padding='same'),\n            nn.ReLU(),\n            nn.Dropout(self.dropout),\n            nn.Conv1d(self.d_model, self.c_out, kernel_size=self.kernel_size, padding='same'),\n        )\n        \n    def anomaly_detection(self, x_enc):\n        \"\"\"\n        Anomaly detection: reconstruct the input sequence.\n        Input: x_enc [batch, seq_len, enc_in]\n        Output: [batch, seq_len, c_out]\n        \"\"\"\n        batch_size, seq_len, n_channels = x_enc.shape\n        \n        # Encoder\n        x = x_enc.permute(0, 2, 1)  # [batch, enc_in, seq_len]\n        x = self.encoder(x)  # [batch, d_model, seq_len]\n        x = x.permute(0, 2, 1)  # [batch, seq_len, d_model]\n        \n        # Add positional encoding\n        x = self.pos_encoder(x)\n        \n        # Transformer encoding\n        x = self.transformer_encoder(x)  # [batch, seq_len, d_model]\n        \n        # Memory addressing\n        # Expand memory to batch dimension\n        memory = self.memory.unsqueeze(0).repeat(batch_size, 1, 1)  # [batch, memory_size, d_model]\n        \n        # Use attention to retrieve from memory\n        # Query: x (seq_len, d_model), Key/Value: memory (memory_size, d_model)\n        attn_output, attn_weights = self.attention(x, memory, memory)  # [batch, seq_len, d_model]\n        \n        # Residual connection\n        x = x + attn_output\n        \n        # Decoder\n        x = x.permute(0, 2, 1)  # [batch, d_model, seq_len]\n        x = self.decoder(x)  # [batch, c_out, seq_len]\n        x = x.permute(0, 2, 1)  # [batch, seq_len, c_out]\n        \n        return x\n\n    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None):\n        if self.task_name == 'anomaly_detection':\n            dec_out = self.anomaly_detection(x_enc)\n            return dec_out\n        return None\n\n\nclass PositionalEncoding(nn.Module):\n    def __init__(self, d_model, max_len=5000):\n        super(PositionalEncoding, self).__init__()\n        pe = torch.zeros(max_len, d_model)\n        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model))\n        pe[:, 0::2] = torch.sin(position * div_term)\n        pe[:, 1::2] = torch.cos(position * div_term)\n        pe = pe.unsqueeze(0)\n        self.register_buffer('pe', pe)\n\n    def forward(self, x):\n        return x + self.pe[:, :x.size(1)]"
7}

1{
2  "op": "replace",
3  "filename": "Time-Series-Library/models/Custom.py",
4  "start_line": 60,
5  "end_line": 70,
6  "content": "        # Memory module\n        self.memory_size = 32\n        self.memory_dim = self.d_model\n        self.memory = nn.Parameter(torch.randn(self.memory_size, self.memory_dim))\n        self.temperature = 1.0\n        self.sparsity_k = 5  # top-k memory items\n        \n        # Learnable scaling for similarity\n        self.sim_scale = nn.Parameter(torch.tensor(10.0))\n        \n        # Decoder: reconstruct from memory-enhanced representation\n        self.decoder = nn.Sequential(\n            nn.Conv1d(self.d_model, self.d_model, kernel_size=self.kernel_size, padding='same'),\n            nn.ReLU(),\n            nn.Dropout(self.dropout),\n            nn.Conv1d(self.d_model, self.c_out, kernel_size=self.kernel_size, padding='same'),\n        )"
7}

1{
2  "op": "replace",
3  "filename": "Time-Series-Library/models/Custom.py",
4  "start_line": 78,
5  "end_line": 113,
6  "content": "    def memory_addressing(self, queries, memory):\n        \"\"\"\n        Sparse memory addressing with top-k selection.\n        queries: [batch, seq_len, d_model]\n        memory: [batch, memory_size, d_model]\n        returns: retrieved memory [batch, seq_len, d_model], weights [batch, seq_len, memory_size]\n        \"\"\"\n        batch_size, seq_len, d_model = queries.shape\n        memory_size = memory.shape[1]\n        \n        # Normalize queries and memory for cosine similarity\n        queries_norm = F.normalize(queries, dim=-1)  # [batch, seq_len, d_model]\n        memory_norm = F.normalize(memory, dim=-1)    # [batch, memory_size, d_model]\n        \n        # Compute cosine similarity\n        sim = torch.bmm(queries_norm, memory_norm.transpose(1, 2))  # [batch, seq_len, memory_size]\n        sim = sim * self.sim_scale  # learnable scaling\n        \n        # Top-k sparsity: keep only top-k similarities, set others to -inf\n        if self.sparsity_k < memory_size:\n            topk_values, topk_indices = torch.topk(sim, k=self.sparsity_k, dim=-1)\n            mask = torch.full_like(sim, float('-inf'))\n            mask.scatter_(-1, topk_indices, topk_values)\n            sim = mask\n        \n        # Softmax with temperature\n        weights = F.softmax(sim / self.temperature, dim=-1)  # [batch, seq_len, memory_size]\n        \n        # Retrieve memory\n        retrieved = torch.bmm(weights, memory)  # [batch, seq_len, d_model]\n        \n        return retrieved, weights\n    \n    def add_noise(self, x):\n        \"\"\"Add Gaussian noise to input during training.\"\"\"\n        if self.training:\n            # Noise standard deviation proportional to data std\n            std = x.std(dim=1, keepdim=True).std(dim=2, keepdim=True).detach()\n            noise = torch.randn_like(x) * std * 0.1\n            x = x + noise\n        return x\n    \n    def anomaly_detection(self, x_enc):\n        \"\"\"\n        Anomaly detection: reconstruct the input sequence.\n        Input: x_enc [batch, seq_len, enc_in]\n        Output: [batch, seq_len, c_out]\n        \"\"\"\n        batch_size, seq_len, n_channels = x_enc.shape\n        \n        # Optionally add noise for denoising\n        x_enc = self.add_noise(x_enc)\n        \n        # Encoder\n        x = x_enc.permute(0, 2, 1)  # [batch, enc_in, seq_len]\n        x = self.encoder(x)  # [batch, d_model, seq_len]\n        x = x.permute(0, 2, 1)  # [batch, seq_len, d_model]\n        \n        # Add positional encoding\n        x = self.pos_encoder(x)\n        \n        # Transformer encoding\n        x = self.transformer_encoder(x)  # [batch, seq_len, d_model]\n        \n        # Memory addressing\n        memory = self.memory.unsqueeze(0).repeat(batch_size, 1, 1)  # [batch, memory_size, d_model]\n        mem_retrieved, attn_weights = self.memory_addressing(x, memory)\n        \n        # Residual connection\n        x = x + mem_retrieved\n        \n        # Decoder\n        x = x.permute(0, 2, 1)  # [batch, d_model, seq_len]\n        x = self.decoder(x)  # [batch, c_out, seq_len]\n        x = x.permute(0, 2, 1)  # [batch, seq_len, c_out]\n        \n        return x"
7}

1{
2  "op": "replace",
3  "filename": "Time-Series-Library/models/Custom.py",
4  "start_line": 60,
5  "end_line": 76,
6  "content": "        # Vector Quantization layer\n        self.codebook_size = 32\n        self.codebook_dim = self.d_model\n        self.commitment_cost = 0.25\n        self.decay = 0.99\n        self.epsilon = 1e-5\n        \n        # Codebook (memory)\n        self.codebook = nn.Parameter(torch.randn(self.codebook_size, self.codebook_dim))\n        self.register_buffer('codebook_ema', torch.zeros(self.codebook_size, self.codebook_dim))\n        self.register_buffer('codebook_ema_count', torch.zeros(self.codebook_size))\n        \n        # Decoder: reconstruct from quantized representation\n        self.decoder = nn.Sequential(\n            nn.Conv1d(self.d_model, self.d_model, kernel_size=self.kernel_size, padding='same'),\n            nn.ReLU(),\n            nn.Dropout(self.dropout),\n            nn.Conv1d(self.d_model, self.c_out, kernel_size=self.kernel_size, padding='same'),\n        )"
7}

1{
2  "op": "replace",
3  "filename": "Time-Series-Library/models/Custom.py",
4  "start_line": 80,
5  "end_line": 156,
6  "content": "    def vector_quantize(self, z):\n        \"\"\"\n        Vector quantization layer.\n        z: [batch, seq_len, d_model]\n        returns: quantized z, codebook indices, commitment loss\n        \"\"\"\n        batch_size, seq_len, d_model = z.shape\n        \n        # Flatten to [batch*seq_len, d_model]\n        z_flat = z.reshape(-1, d_model)\n        \n        # Compute distances to codebook vectors\n        distances = torch.cdist(z_flat, self.codebook, p=2)  # [batch*seq_len, codebook_size]\n        \n        # Find nearest codebook indices\n        indices = torch.argmin(distances, dim=-1)  # [batch*seq_len]\n        \n        # Quantized vectors\n        quantized_flat = self.codebook[indices]  # [batch*seq_len, d_model]\n        \n        # Reshape back\n        quantized = quantized_flat.reshape(batch_size, seq_len, d_model)\n        indices = indices.reshape(batch_size, seq_len)\n        \n        # EMA update of codebook (during training)\n        if self.training:\n            # Use exponential moving average to update codebook vectors\n            with torch.no_grad():\n                # Update EMA counts\n                indices_onehot = F.one_hot(indices, self.codebook_size).float()  # [batch, seq_len, codebook_size]\n                indices_onehot_flat = indices_onehot.reshape(-1, self.codebook_size)  # [batch*seq_len, codebook_size]\n                counts = indices_onehot_flat.sum(0)  # [codebook_size]\n                \n                # Update EMA\n                self.codebook_ema_count.mul_(self.decay).add_(counts, alpha=1 - self.decay)\n                \n                # Compute sum of z_flat per codebook index\n                z_sum = torch.matmul(indices_onehot_flat.t(), z_flat)  # [codebook_size, d_model]\n                self.codebook_ema.mul_(self.decay).add_(z_sum, alpha=1 - self.decay)\n                \n                # Normalize to get updated codebook\n                n = self.codebook_ema_count.unsqueeze(1)\n                n = torch.max(n, torch.tensor(self.epsilon))\n                updated_codebook = self.codebook_ema / n\n                \n                # Update codebook parameters (with gradient bypass)\n                self.codebook.data = updated_codebook\n        \n        # Straight-through gradient estimator: copy gradient from quantized to z\n        quantized = z + (quantized - z).detach()\n        \n        # Commitment loss (optional) - we could compute but cannot add to loss directly\n        # Instead we can add a small penalty to the quantized representation\n        commitment_loss = F.mse_loss(z.detach(), quantized)\n        # We'll not use this loss because we cannot add it to total loss\n        \n        return quantized, indices\n    \n    def add_noise(self, x):\n        \"\"\"Add Gaussian noise to input during training.\"\"\"\n        if self.training:\n            # Noise standard deviation proportional to data std\n            std = x.std(dim=1, keepdim=True).std(dim=2, keepdim=True).detach()\n            noise = torch.randn_like(x) * std * 0.1\n            x = x + noise\n        return x\n    \n    def anomaly_detection(self, x_enc):\n        \"\"\"\n        Anomaly detection: reconstruct the input sequence.\n        Input: x_enc [batch, seq_len, enc_in]\n        Output: [batch, seq_len, c_out]\n        \"\"\"\n        batch_size, seq_len, n_channels = x_enc.shape\n        \n        # Optionally add noise for denoising\n        x_enc = self.add_noise(x_enc)\n        \n        # Encoder\n        x = x_enc.permute(0, 2, 1)  # [batch, enc_in, seq_len]\n        x = self.encoder(x)  # [batch, d_model, seq_len]\n        x = x.permute(0, 2, 1)  # [batch, seq_len, d_model]\n        \n        # Add positional encoding\n        x = self.pos_encoder(x)\n        \n        # Transformer encoding\n        x = self.transformer_encoder(x)  # [batch, seq_len, d_model]\n        \n        # Vector quantization\n        x_quantized, indices = self.vector_quantize(x)\n        \n        # Residual connection (skip connection around quantization)\n        x = x + x_quantized\n        \n        # Decoder\n        x = x.permute(0, 2, 1)  # [batch, d_model, seq_len]\n        x = self.decoder(x)  # [batch, c_out, seq_len]\n        x = x.permute(0, 2, 1)  # [batch, seq_len, c_out]\n        \n        return x"
7}

1{
2  "op": "replace",
3  "filename": "Time-Series-Library/models/Custom.py",
4  "start_line": 29,
5  "end_line": 35,
6  "content": "        # Hyperparameters\n        self.d_model = 64\n        self.num_layers = 2\n        self.kernel_size = 3\n        self.dropout = 0.1"
7}

1{
2  "op": "replace",
3  "filename": "Time-Series-Library/models/Custom.py",
4  "start_line": 170,
5  "end_line": 172,
6  "content": "          # Use quantized representation as input to decoder\n          x = x_quantized"
7}

1{
2  "op": "replace",
3  "filename": "Time-Series-Library/models/Custom.py",
4  "start_line": 169,
5  "end_line": 172,
6  "content": "          # Use quantized representation as input to decoder\n          x = x_quantized"
7}