Upload 4 files

TimesNet.py

@@ -0,0 +1,418 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.fft
+import numpy as np
+# Basit embedding ve conv blocks - layers klasörü olmadan
+class DataEmbedding(nn.Module):
+    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1, seq_len=6000):
+        super(DataEmbedding, self).__init__()
+        self.c_in = c_in
+        self.d_model = d_model
+        self.embed_type = embed_type
+        self.freq = freq
+        self.seq_len = seq_len
+        
+        # Basit linear embedding
+        self.value_embedding = nn.Linear(c_in, d_model)
+        # Position embedding'i seq_len'e göre oluştur
+        self.position_embedding = nn.Parameter(torch.randn(1, seq_len, d_model))
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x, x_mark):
+        x = self.value_embedding(x)
+        
+        # Position embedding'i input boyutuna göre crop et
+        # seq_len'e göre oluşturulduğu için genelde uyumlu olacak
+        if x.size(1) <= self.position_embedding.size(1):
+            x = x + self.position_embedding[:, :x.size(1), :]
+        else:
+            # Eğer input daha büyükse, position embedding'i extend et
+            x = x + self.position_embedding
+            remaining_length = x.size(1) - self.position_embedding.size(1)
+            if remaining_length > 0:
+                # Sinusoidal position encoding ekle
+                pos_encoding = self._get_sinusoidal_encoding(remaining_length, self.d_model)
+                pos_encoding = pos_encoding.unsqueeze(0).to(x.device)
+                x[:, self.position_embedding.size(1):, :] += pos_encoding
+        
+        return self.dropout(x)
+    
+    def _get_sinusoidal_encoding(self, length, d_model):
+        """Sinusoidal position encoding oluştur"""
+        position = torch.arange(length).unsqueeze(1).float()
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(np.log(10000.0) / d_model))
+        
+        pos_encoding = torch.zeros(length, d_model)
+        pos_encoding[:, 0::2] = torch.sin(position * div_term)
+        pos_encoding[:, 1::2] = torch.cos(position * div_term)
+        
+        return pos_encoding
+
+class Inception_Block_V1(nn.Module):
+    def __init__(self, in_channels, out_channels, num_kernels=6, init_weight=True):
+        super(Inception_Block_V1, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_kernels = num_kernels
+        kernels = []
+        for i in range(self.num_kernels):
+            kernels.append(nn.Conv2d(in_channels, out_channels, kernel_size=2 * i + 1, padding=i))
+        self.kernels = nn.ModuleList(kernels)
+        if init_weight:
+            self._initialize_weights()
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        res_list = []
+        for i, kernel in enumerate(self.kernels):
+            res_list.append(kernel(x))
+        res = torch.stack(res_list, dim=-1).mean(-1)
+        return res
+
+
+def FFT_for_Period(x, k=2):
+    # [B, T, C]
+    xf = torch.fft.rfft(x, dim=1)
+    # find period by amplitudes
+    frequency_list = abs(xf).mean(0).mean(-1)
+    frequency_list[0] = 0
+    _, top_list = torch.topk(frequency_list, k)
+    top_list = top_list.detach().cpu().numpy()
+    period = x.shape[1] // top_list
+    return period, abs(xf).mean(-1)[:, top_list]
+
+
+class TimesBlock(nn.Module):
+    def __init__(self, configs):
+        super(TimesBlock, self).__init__()
+        self.seq_len = configs.seq_len
+        self.pred_len = configs.pred_len
+        self.k = configs.top_k
+        # parameter-efficient design
+        self.conv = nn.Sequential(
+            Inception_Block_V1(configs.d_model, configs.d_ff,
+                               num_kernels=configs.num_kernels),
+            nn.GELU(),
+            Inception_Block_V1(configs.d_ff, configs.d_model,
+                               num_kernels=configs.num_kernels)
+        )
+
+    def forward(self, x):
+        B, T, N = x.size() #B: batch size  T: length of time series  N:number of features
+        period_list, period_weight = FFT_for_Period(x, self.k)
+
+        res = []
+        for i in range(self.k):
+            period = period_list[i]
+            # padding
+            if (self.seq_len + self.pred_len) % period != 0:
+                length = (
+                                 ((self.seq_len + self.pred_len) // period) + 1) * period
+                padding = torch.zeros([x.shape[0], (length - (self.seq_len + self.pred_len)), x.shape[2]]).to(x.device)
+                out = torch.cat([x, padding], dim=1)
+            else:
+                length = (self.seq_len + self.pred_len)
+                out = x
+            # reshape
+            out = out.reshape(B, length // period, period,
+                              N).permute(0, 3, 1, 2).contiguous()
+            # 2D conv: from 1d Variation to 2d Variation
+            out = self.conv(out)
+            # reshape back
+            out = out.permute(0, 2, 3, 1).reshape(B, -1, N)
+            res.append(out[:, :(self.seq_len + self.pred_len), :])
+        res = torch.stack(res, dim=-1)
+        # adaptive aggregation
+        period_weight = F.softmax(period_weight, dim=1)
+        period_weight = period_weight.unsqueeze(
+            1).unsqueeze(1).repeat(1, T, N, 1)
+        res = torch.sum(res * period_weight, -1)
+        # residual connection
+        res = res + x
+        return res
+
+
+class Model(nn.Module):
+    """
+    Paper link: https://openreview.net/pdf?id=ju_Uqw384Oq
+    """
+
+    def __init__(self, configs):
+        super(Model, self).__init__()
+        self.configs = configs
+        self.task_name = configs.task_name
+        self.seq_len = configs.seq_len
+        self.label_len = configs.label_len
+        self.pred_len = configs.pred_len
+        self.model = nn.ModuleList([TimesBlock(configs)
+                                    for _ in range(configs.e_layers)])
+        self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                           configs.dropout, configs.seq_len)
+        self.layer = configs.e_layers
+        self.layer_norm = nn.LayerNorm(configs.d_model)
+        if self.task_name == 'long_term_forecast' or self.task_name == 'short_term_forecast':
+            self.predict_linear = nn.Linear(
+                self.seq_len, self.pred_len + self.seq_len)
+            self.projection = nn.Linear(
+                configs.d_model, configs.c_out, bias=True)
+        if self.task_name == 'imputation' or self.task_name == 'anomaly_detection':
+            self.projection = nn.Linear(
+                configs.d_model, configs.c_out, bias=True)
+            
+        # Transfer learning için P-S prediction heads (sadece gerektiğinde eklenir)
+        if hasattr(configs, 'use_ps_heads') and configs.use_ps_heads:
+            # Skip attention for memory efficiency - use only pooling
+            
+            # Multi-scale feature extraction (reduced sizes for memory)
+            self.multi_scale_pools = nn.ModuleList([
+                nn.AdaptiveAvgPool1d(16),   # Local patterns (reduced)
+                nn.AdaptiveAvgPool1d(4),    # Medium patterns
+                nn.AdaptiveAvgPool1d(1),    # Global patterns
+            ])
+            
+            # Feature fusion - calculate exact dimension
+            # Pool sizes: 16 + 4 + 1 = 21, so total dim = d_model * 21
+            fusion_dim = configs.d_model * (16 + 4 + 1)  # Exact calculation
+            self.feature_fusion = nn.Sequential(
+                nn.Linear(fusion_dim, configs.d_model),
+                nn.ReLU(),
+                nn.Dropout(configs.dropout)
+            )
+            
+            # Separate P and S regression heads
+            self.p_regression_head = nn.Sequential(
+                nn.Linear(configs.d_model, 128),
+                nn.ReLU(),
+                nn.Dropout(configs.dropout),
+                nn.Linear(128, 64),
+                nn.ReLU(),
+                nn.Dropout(configs.dropout),
+                nn.Linear(64, 1)  # P time only
+            )
+            
+            self.s_regression_head = nn.Sequential(
+                nn.Linear(configs.d_model, 128),
+                nn.ReLU(),
+                nn.Dropout(configs.dropout),
+                nn.Linear(128, 64),
+                nn.ReLU(),
+                nn.Dropout(configs.dropout),
+                nn.Linear(64, 1)  # S time only
+            )
+            
+            # Separate P and S classification heads
+            self.p_classification_head = nn.Sequential(
+                nn.Linear(configs.d_model, 64),
+                nn.ReLU(),
+                nn.Dropout(configs.dropout),
+                nn.Linear(64, 32),
+                nn.ReLU(),
+                nn.Dropout(configs.dropout),
+                nn.Linear(32, 1),  # P exists/not
+                nn.Sigmoid()
+            )
+            
+            self.s_classification_head = nn.Sequential(
+                nn.Linear(configs.d_model, 64),
+                nn.ReLU(),
+                nn.Dropout(configs.dropout),
+                nn.Linear(64, 32),
+                nn.ReLU(),
+                nn.Dropout(configs.dropout),
+                nn.Linear(32, 1),  # S exists/not
+                nn.Sigmoid()
+            )
+        if self.task_name == 'classification':
+            self.act = F.gelu
+            self.dropout = nn.Dropout(configs.dropout)
+            self.projection = nn.Linear(
+                configs.d_model * configs.seq_len, configs.num_class)
+
+    def anomaly_detection(self, x_enc):
+        # Transfer learning için P-S heads varsa - SADECE ONLARI KULLAN
+        if hasattr(self, 'p_regression_head'):
+            # Normalization from Non-stationary Transformer
+            means = x_enc.mean(1, keepdim=True).detach()
+            x_enc = x_enc - means
+            stdev = torch.sqrt(
+                torch.var(x_enc, dim=1, keepdim=True, unbiased=False) + 1e-5)
+            x_enc /= stdev
+
+            # embedding
+            enc_out = self.enc_embedding(x_enc, None)  # [B,T,C]
+            # TimesNet
+            for i in range(self.layer):
+                enc_out = self.layer_norm(self.model[i](enc_out))
+            
+            # Skip attention for memory - use direct multi-scale pooling
+            # Multi-scale feature extraction directly on TimesNet output
+            enc_out_transposed = enc_out.permute(0, 2, 1)  # (B, d_model, T)
+            multi_scale_features = []
+            
+            # Manual pooling for large sequences to avoid CUDA memory issues
+            pool_sizes = [16, 4, 1]  # Target pool sizes
+            for i, target_size in enumerate(pool_sizes):
+                T = enc_out_transposed.size(2)  # Sequence length
+                
+                if T >= 8000:  # Very large - use manual avg pooling
+                    # Manual average pooling
+                    window_size = T // target_size
+                    if window_size > 0:
+                        # Reshape and average
+                        # (B, d_model, T) -> (B, d_model, target_size, window_size)
+                        trimmed_T = (T // window_size) * window_size
+                        trimmed = enc_out_transposed[:, :, :trimmed_T]
+                        reshaped = trimmed.view(trimmed.size(0), trimmed.size(1), target_size, window_size)
+                        pooled = reshaped.mean(dim=3)  # Average over window
+                    else:
+                        # Fallback: simple reshape
+                        pooled = enc_out_transposed[:, :, :target_size] if T >= target_size else enc_out_transposed
+                else:
+                    # Use normal adaptive pooling for smaller sequences
+                    pool = self.multi_scale_pools[i]
+                    pooled = pool(enc_out_transposed)  # (B, d_model, pool_size)
+                
+                flattened = pooled.flatten(1)  # (B, d_model * pool_size)
+                multi_scale_features.append(flattened)
+            
+            # Concatenate multi-scale features
+            fused_features = torch.cat(multi_scale_features, dim=1)  # (B, d_model * 3)
+            
+            # Feature fusion
+            final_features = self.feature_fusion(fused_features)  # (B, d_model)
+            
+            # Separate P and S predictions
+            p_time = self.p_regression_head(final_features)  # (B, 1)
+            s_time = self.s_regression_head(final_features)  # (B, 1)
+            ps_times = torch.cat([p_time, s_time], dim=1)  # (B, 2)
+            
+            # Separate P and S classifications
+            p_class = self.p_classification_head(final_features)  # (B, 1)
+            s_class = self.s_classification_head(final_features)  # (B, 1)
+            ps_classification = torch.cat([p_class, s_class], dim=1)  # (B, 2)
+            
+            return ps_times, ps_classification
+        else:
+            # Orijinal anomaly detection (reconstruction)
+            # Normalization from Non-stationary Transformer
+            means = x_enc.mean(1, keepdim=True).detach()
+            x_enc = x_enc - means
+            stdev = torch.sqrt(
+                torch.var(x_enc, dim=1, keepdim=True, unbiased=False) + 1e-5)
+            x_enc /= stdev
+
+            # embedding
+            enc_out = self.enc_embedding(x_enc, None)  # [B,T,C]
+            # TimesNet
+            for i in range(self.layer):
+                enc_out = self.layer_norm(self.model[i](enc_out))
+            # porject back
+            dec_out = self.projection(enc_out)
+
+            # De-Normalization from Non-stationary Transformer
+            dec_out = dec_out * \
+                      (stdev[:, 0, :].unsqueeze(1).repeat(
+                          1, self.pred_len + self.seq_len, 1))
+            dec_out = dec_out + \
+                      (means[:, 0, :].unsqueeze(1).repeat(
+                          1, self.pred_len + self.seq_len, 1))
+            return dec_out
+
+    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None):
+        if self.task_name == 'long_term_forecast' or self.task_name == 'short_term_forecast':
+            dec_out = self.forecast(x_enc, x_mark_enc, x_dec, x_mark_dec)
+            return dec_out[:, -self.pred_len:, :]  # [B, L, D]
+        if self.task_name == 'imputation':
+            dec_out = self.imputation(
+                x_enc, x_mark_enc, x_dec, x_mark_dec, mask)
+            return dec_out  # [B, L, D]
+        if self.task_name == 'anomaly_detection':
+            result = self.anomaly_detection(x_enc)
+            return result  # [B, L, D] veya [B, L, D], [B, 2], [B, 1]
+        if self.task_name == 'classification':
+            dec_out = self.classification(x_enc, x_mark_enc)
+            return dec_out  # [B, N]
+        return None
+
+    def forecast(self, x_enc, x_mark_enc, x_dec, x_mark_dec):
+        # Normalization from Non-stationary Transformer
+        means = x_enc.mean(1, keepdim=True).detach()
+        x_enc = x_enc - means
+        stdev = torch.sqrt(
+            torch.var(x_enc, dim=1, keepdim=True, unbiased=False) + 1e-5)
+        x_enc /= stdev
+
+        # embedding
+        enc_out = self.enc_embedding(x_enc, x_mark_enc)  # [B,T,C]
+        enc_out = self.predict_linear(enc_out.permute(0, 2, 1)).permute(
+            0, 2, 1)  # align temporal dimension
+        # TimesNet
+        for i in range(self.layer):
+            enc_out = self.layer_norm(self.model[i](enc_out))
+        # porject back
+        dec_out = self.projection(enc_out)
+
+        # De-Normalization from Non-stationary Transformer
+        dec_out = dec_out * \
+                  (stdev[:, 0, :].unsqueeze(1).repeat(
+                      1, self.pred_len + self.seq_len, 1))
+        dec_out = dec_out + \
+                  (means[:, 0, :].unsqueeze(1).repeat(
+                      1, self.pred_len + self.seq_len, 1))
+        return dec_out
+
+    def imputation(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask):
+        # Normalization from Non-stationary Transformer
+        means = torch.sum(x_enc, dim=1) / torch.sum(mask == 1, dim=1)
+        means = means.unsqueeze(1).detach()
+        x_enc = x_enc - means
+        x_enc = x_enc.masked_fill(mask == 0, 0)
+        stdev = torch.sqrt(torch.sum(x_enc * x_enc, dim=1) /
+                           torch.sum(mask == 1, dim=1) + 1e-5)
+        stdev = stdev.unsqueeze(1).detach()
+        x_enc /= stdev
+
+        # embedding
+        enc_out = self.enc_embedding(x_enc, x_mark_enc)  # [B,T,C]
+        # TimesNet
+        for i in range(self.layer):
+            enc_out = self.layer_norm(self.model[i](enc_out))
+        # porject back
+        dec_out = self.projection(enc_out)
+
+        # De-Normalization from Non-stationary Transformer
+        dec_out = dec_out * \
+                  (stdev[:, 0, :].unsqueeze(1).repeat(
+                      1, self.pred_len + self.seq_len, 1))
+        dec_out = dec_out + \
+                  (means[:, 0, :].unsqueeze(1).repeat(
+                      1, self.pred_len + self.seq_len, 1))
+        return dec_out
+
+
+
+    def classification(self, x_enc, x_mark_enc):
+        # embedding
+        enc_out = self.enc_embedding(x_enc, None)  # [B,T,C]
+        # TimesNet
+        for i in range(self.layer):
+            enc_out = self.layer_norm(self.model[i](enc_out))
+
+        # Output
+        # the output transformer encoder/decoder embeddings don't include non-linearity
+        output = self.act(enc_out)
+        output = self.dropout(output)
+        # zero-out padding embeddings
+        output = output * x_mark_enc.unsqueeze(-1)
+        # (batch_size, seq_length * d_model)
+        output = output.reshape(output.shape[0], -1)
+        output = self.projection(output)  # (batch_size, num_classes)
+        return output
+

TimesNet_PointCloud.py

@@ -0,0 +1,213 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+try:
+    from .TimesNet import DataEmbedding
+except Exception:
+    from TimesNet import DataEmbedding
+
+
+class _BlockConfig:
+    def __init__(self, seq_len: int, pred_len: int, d_model: int, d_ff: int, num_kernels: int, top_k: int = 2, num_stations: int = 0):
+        self.seq_len = seq_len
+        self.pred_len = pred_len
+        self.d_model = d_model
+        self.d_ff = d_ff
+        self.num_kernels = num_kernels
+        self.top_k = top_k
+        self.num_stations = num_stations
+
+
+class Inception_Block_V1(nn.Module):
+    def __init__(self, in_channels, out_channels, num_kernels=6, init_weight=True):
+        super(Inception_Block_V1, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_kernels = num_kernels
+        kernels = []
+        for i in range(self.num_kernels):
+            kernels.append(nn.Conv2d(in_channels, out_channels, kernel_size=2 * i + 1, padding=i))
+        self.kernels = nn.ModuleList(kernels)
+        if init_weight:
+            self._initialize_weights()
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        res_list = []
+        for i, kernel in enumerate(self.kernels):
+            res_list.append(kernel(x))
+        res = torch.stack(res_list, dim=-1).mean(-1)
+        return res
+
+
+def FFT_for_Period(x, k=2):
+    # [B, T, C]
+    xf = torch.fft.rfft(x, dim=1)
+    # find period by amplitudes
+    frequency_list = abs(xf).mean(0).mean(-1)
+    frequency_list[0] = 0
+    _, top_list = torch.topk(frequency_list, k)
+    top_list = top_list.detach().cpu().numpy()
+    period = x.shape[1] // top_list
+    return period, abs(xf).mean(-1)[:, top_list]
+
+
+class TimesBlockStationCond(nn.Module):
+    """TimesBlock with station ID conditioning (one-hot encoded as 1 channel)."""
+    def __init__(self, configs):
+        super(TimesBlockStationCond, self).__init__()
+        self.seq_len = configs.seq_len
+        self.pred_len = configs.pred_len
+        self.k = configs.top_k
+        self.num_stations = getattr(configs, 'num_stations', 0)
+        
+        # Station ID embedding: maps station ID to d_model dimension
+        # This provides richer conditioning information than a single scalar
+        if self.num_stations > 0:
+            self.station_embedding = nn.Embedding(self.num_stations, configs.d_model)
+            # Initialize with small random values
+            nn.init.normal_(self.station_embedding.weight, mean=0.0, std=0.02)
+        
+        # Inception blocks
+        self.conv = nn.Sequential(
+            Inception_Block_V1(configs.d_model, configs.d_ff,
+                               num_kernels=configs.num_kernels),
+            nn.GELU(),
+            Inception_Block_V1(configs.d_ff, configs.d_model,
+                               num_kernels=configs.num_kernels)
+        )
+
+    def forward(self, x, station_ids: torch.Tensor = None):
+        """
+        Args:
+            x: (B, T, N) input features
+            station_ids: (B,) LongTensor of station IDs (0 to num_stations-1)
+        """
+        B, T, N = x.size()
+        period_list, period_weight = FFT_for_Period(x, self.k)
+
+        res = []
+        for i in range(self.k):
+            period = period_list[i]
+            # padding
+            if (self.seq_len + self.pred_len) % period != 0:
+                length = (((self.seq_len + self.pred_len) // period) + 1) * period
+                padding = torch.zeros([x.shape[0], (length - (self.seq_len + self.pred_len)), x.shape[2]]).to(x.device)
+                out = torch.cat([x, padding], dim=1)
+            else:
+                length = (self.seq_len + self.pred_len)
+                out = x
+            
+            # reshape to 2D: (B, N, H, W)
+            out = out.reshape(B, length // period, period, N).permute(0, 3, 1, 2).contiguous()
+            
+            # Inject station ID conditioning via embedding addition
+            # This provides richer conditioning (d_model dimensions) compared to single scalar
+            if station_ids is not None and self.num_stations > 0:
+                # Get station embeddings: (B, d_model)
+                station_ids_flat = station_ids.view(B)
+                station_emb = self.station_embedding(station_ids_flat)  # (B, d_model)
+                
+                # out shape: (B, d_model, H, W)
+                # Expand station embedding to spatial dimensions: (B, d_model, H, W)
+                H = out.size(2)
+                W = out.size(3)
+                station_emb_spatial = station_emb.view(B, N, 1, 1).expand(-1, -1, H, W)
+                
+                # Add station embedding to features (element-wise addition)
+                # This allows the model to learn station-specific feature modifications
+                out = out + station_emb_spatial
+            
+            # 2D conv: from 1d Variation to 2d Variation
+            out = self.conv(out)
+            
+            # reshape back
+            out = out.permute(0, 2, 3, 1).reshape(B, -1, N)
+            res.append(out[:, :(self.seq_len + self.pred_len), :])
+        
+        res = torch.stack(res, dim=-1)
+        # adaptive aggregation
+        period_weight = F.softmax(period_weight, dim=1)
+        period_weight = period_weight.unsqueeze(1).unsqueeze(1).repeat(1, T, N, 1)
+        res = torch.sum(res * period_weight, -1)
+        
+        # residual connection
+        res = res + x
+        return res
+
+
+class TimesNetPointCloud(nn.Module):
+    """TimesNet reconstruction with exposed encode/project methods for point-cloud mixing."""
+    def __init__(self, configs):
+        super().__init__()
+        self.configs = configs
+        self.seq_len = configs.seq_len
+        self.pred_len = getattr(configs, 'pred_len', 0)
+        self.top_k = configs.top_k
+        self.d_model = configs.d_model
+        self.d_ff = configs.d_ff
+        self.num_kernels = configs.num_kernels
+        self.e_layers = configs.e_layers
+        self.dropout = configs.dropout
+        self.c_out = configs.c_out
+
+        self.num_stations = getattr(configs, 'num_stations', 0)
+
+        self.enc_embedding = DataEmbedding(configs.enc_in, self.d_model, configs.embed, configs.freq,
+                                           configs.dropout, configs.seq_len)
+        self.model = nn.ModuleList([
+            TimesBlockStationCond(_BlockConfig(self.seq_len, 0, self.d_model, self.d_ff, 
+                                              self.num_kernels, self.top_k, self.num_stations))
+            for _ in range(self.e_layers)
+        ])
+        self.layer = self.e_layers
+        self.layer_norm = nn.LayerNorm(self.d_model)
+        self.projection = nn.Linear(self.d_model, self.c_out, bias=True)
+
+    def encode_features_for_reconstruction(self, x_enc: torch.Tensor, station_ids: torch.Tensor = None):
+        """
+        Encode input with optional station ID conditioning.
+        
+        Args:
+            x_enc: (B, T, C) input signal
+            station_ids: (B,) LongTensor of station IDs (0 to num_stations-1), optional
+        """
+        means = x_enc.mean(1, keepdim=True).detach()
+        x_norm = x_enc - means
+        stdev = torch.sqrt(torch.var(x_norm, dim=1, keepdim=True, unbiased=False) + 1e-5)
+        x_norm = x_norm / stdev
+        enc_out = self.enc_embedding(x_norm, None)
+        for i in range(self.layer):
+            enc_out = self.layer_norm(self.model[i](enc_out, station_ids))
+        return enc_out, means, stdev
+
+    def project_features_for_reconstruction(self, enc_out: torch.Tensor, means: torch.Tensor, stdev: torch.Tensor):
+        dec_out = self.projection(enc_out)
+        dec_out = dec_out * (stdev[:, 0, :].unsqueeze(1).repeat(1, self.pred_len + self.seq_len, 1))
+        dec_out = dec_out + (means[:, 0, :].unsqueeze(1).repeat(1, self.pred_len + self.seq_len, 1))
+        return dec_out
+
+    def anomaly_detection(self, x_enc: torch.Tensor, station_ids: torch.Tensor = None):
+        """Full reconstruction pass with optional station ID conditioning."""
+        enc_out, means, stdev = self.encode_features_for_reconstruction(x_enc, station_ids)
+        return self.project_features_for_reconstruction(enc_out, means, stdev)
+
+    def forward(self, x_enc, station_ids=None, x_mark_enc=None, x_dec=None, x_mark_dec=None, mask=None):
+        """
+        Forward pass compatible with anomaly_detection task.
+        
+        Args:
+            x_enc: (B, T, C) input signal
+            station_ids: (B,) LongTensor of station IDs, optional
+        """
+        return self.anomaly_detection(x_enc, station_ids)
+
+

app.py

@@ -0,0 +1,236 @@
+import gradio as gr
+import numpy as np
+import matplotlib
+matplotlib.use('Agg')  # Use non-interactive backend
+import matplotlib.pyplot as plt
+from PIL import Image
+import io
+import os
+import torch
+
+# Import model and generation functions
+try:
+    from TimesNet_PointCloud import TimesNetPointCloud
+    from generate_samples_git import SimpleArgs, generate_samples_from_latent_bank
+except ImportError:
+    # Fallback for local imports
+    import sys
+    sys.path.insert(0, '.')
+    from TimesNet_PointCloud import TimesNetPointCloud
+    try:
+        from generate_samples_git import SimpleArgs, generate_samples_from_latent_bank
+    except:
+        # If generate_samples_git doesn't exist, use generate_samples
+        from generate_samples import FineTuneArgs as SimpleArgs, generate_samples_from_latent_bank
+
+def load_model(checkpoint_path, args):
+    """Load pre-trained TimesNet-PointCloud model (matching generate_samples_git.py)."""
+    # Create model config
+    class ModelConfig:
+        def __init__(self, args):
+            self.seq_len = args.seq_len
+            self.pred_len = 0
+            self.enc_in = 3
+            self.c_out = 3
+            self.d_model = args.d_model
+            self.d_ff = args.d_ff
+            self.num_kernels = args.num_kernels
+            self.top_k = args.top_k
+            self.e_layers = args.e_layers
+            self.d_layers = args.d_layers
+            self.dropout = args.dropout
+            self.embed = 'timeF'
+            self.freq = 'h'
+            self.latent_dim = getattr(args, 'latent_dim', 256)
+            # num_stations is needed for station conditioning
+            # Try to get from checkpoint, or use default
+            self.num_stations = getattr(args, 'num_stations', 705)
+    
+    config = ModelConfig(args)
+    model = TimesNetPointCloud(config)
+    
+    # Load checkpoint
+    checkpoint = torch.load(checkpoint_path, map_location='cpu')
+    if 'model_state_dict' in checkpoint:
+        model.load_state_dict(checkpoint['model_state_dict'])
+        # Try to get num_stations from checkpoint
+        if 'num_stations' in checkpoint:
+            config.num_stations = checkpoint['num_stations']
+    else:
+        model.load_state_dict(checkpoint)
+    
+    model.eval()
+    if args.use_gpu:
+        model = model.cuda()
+    
+    print(f"[INFO] Model loaded successfully from {checkpoint_path}")
+    return model
+
+# Configuration - can be set via environment variables for HF Space
+PHASE1_MODEL_CHECKPOINT_PATH = os.getenv('PHASE1_MODEL_CHECKPOINT_PATH', './checkpoints/timesnet_pointcloud_phase1_final.pth')
+LATENT_BANK_PATH = os.getenv('LATENT_BANK_PATH', './latent_bank_station_cond.npz')
+ENCODER_STD_PATH = os.getenv('ENCODER_STD_PATH', './pcgen_stats/encoder_feature_std.npy')
+
+# Test stations (5 unseen stations)
+TEST_STATIONS = ['0205', '1716', '2020', '3130', '4628']
+
+# Initialize model and args (loaded once at startup)
+model = None
+args = None
+encoder_std = None
+
+def initialize_model():
+    """Load Phase 1 model and encoder_std once at startup."""
+    global model, args, encoder_std
+    
+    if model is not None:
+        return True  # Already initialized
+    
+    print("[INFO] Initializing TimesNet-Gen Phase 1 model...")
+    
+    # Create args (matching generate_samples_git.py)
+    args = SimpleArgs()
+    
+    # Load Phase 1 model
+    if not os.path.exists(PHASE1_MODEL_CHECKPOINT_PATH):
+        print(f"[ERROR] Phase 1 model checkpoint not found: {PHASE1_MODEL_CHECKPOINT_PATH}")
+        print("[ERROR] Please set PHASE1_MODEL_CHECKPOINT_PATH environment variable")
+        return False
+    
+    model = load_model(PHASE1_MODEL_CHECKPOINT_PATH, args)
+    
+    # Check if latent bank exists
+    if not os.path.exists(LATENT_BANK_PATH):
+        print(f"[ERROR] Latent bank not found: {LATENT_BANK_PATH}")
+        print("[ERROR] Please set LATENT_BANK_PATH environment variable or create latent bank first")
+        return False
+    
+    print(f"[INFO] Using latent bank: {LATENT_BANK_PATH}")
+    
+    # Load encoder_std (optional, only for fine-tuning, not for generation)
+    if os.path.exists(ENCODER_STD_PATH):
+        encoder_std = np.load(ENCODER_STD_PATH)
+        print(f"[INFO] Loaded encoder_std from {ENCODER_STD_PATH} (shape: {encoder_std.shape})")
+        print(f"[INFO] encoder_std loaded (used only for fine-tuning, NOT for generation)")
+    else:
+        print(f"[INFO] No encoder_std found (not needed for generation, only for fine-tuning)")
+        encoder_std = None
+    
+    print("[INFO] ✓ Phase 1 model initialized successfully!")
+    return True
+
+# Initialize on import
+try:
+    initialize_model()
+except Exception as e:
+    print(f"[ERROR] Failed to initialize model: {e}")
+    import traceback
+    traceback.print_exc()
+    print("[WARN] App will run in dummy mode")
+
+def generate_seismic_data(station_id_str, num_samples):
+    """Generate seismic signals using Phase 1 model and pre-computed latent bank."""
+    global model, args, encoder_std
+    
+    # Check if model is loaded
+    if model is None:
+        print("[ERROR] Model not initialized! Attempting to initialize...")
+        if not initialize_model():
+            # Fallback to dummy generation
+            print("[WARN] Using dummy generation as fallback")
+            generated_signals = np.random.randn(num_samples, 3, 6000) * 0.1 + np.sin(np.linspace(0, 100, 6000))
+        else:
+            # Retry generation with real model
+            return generate_seismic_data(station_id_str, num_samples)
+    
+    # Generate samples from pre-computed latent bank (matching generate_samples_git.py)
+    try:
+        print(f"[INFO] Generating {num_samples} samples for station {station_id_str} from latent bank...")
+        
+        # Note: encoder_std is passed but NOT used during generation in generate_samples_git.py
+        # (see line 313-317 in generate_samples_git.py - noise is NOT added during generation)
+        generated_signals, _ = generate_samples_from_latent_bank(
+            model, LATENT_BANK_PATH, station_id_str, num_samples, args, encoder_std
+        )
+        
+        if generated_signals is None:
+            print(f"[ERROR] Failed to generate samples for station {station_id_str}")
+            generated_signals = None
+        else:
+            print(f"[INFO] ✓ Generated {len(generated_signals)} samples successfully")
+    except Exception as e:
+        print(f"[ERROR] Generation failed: {e}")
+        import traceback
+        traceback.print_exc()
+        generated_signals = None
+    
+    # Handle case where generation failed
+    if generated_signals is None:
+        # Create error message plot
+        fig, ax = plt.subplots(1, 1, figsize=(8, 4))
+        ax.text(0.5, 0.5, f'Error: Could not generate samples\nfor station {station_id_str}', 
+                ha='center', va='center', fontsize=14, transform=ax.transAxes)
+        ax.set_xticks([])
+        ax.set_yticks([])
+        plt.tight_layout()
+    else:
+        # generated_signals shape: (num_samples, 3, 6000)
+        num_plots = min(len(generated_signals), 5)  # Show up to 5 samples
+        
+        if num_plots == 1:
+            # Single plot
+            fig, axes = plt.subplots(3, 1, figsize=(10, 6), sharex=True)
+            channel_names = ['E-W', 'N-S', 'U-D']
+            for ch, ax in enumerate(axes):
+                ax.plot(generated_signals[0, ch, :], linewidth=0.8)
+                ax.set_ylabel(channel_names[ch], fontweight='bold')
+                ax.grid(True, alpha=0.3)
+            axes[-1].set_xlabel('Time Steps', fontweight='bold')
+            fig.suptitle(f'Generated Sample for Station {station_id_str}', fontsize=12, fontweight='bold')
+            plt.tight_layout()
+        else:
+            # Multiple plots in a grid
+            fig, axes = plt.subplots(num_plots, 3, figsize=(12, 2*num_plots), sharex=True)
+            channel_names = ['E-W', 'N-S', 'U-D']
+            
+            for i in range(num_plots):
+                for ch in range(3):
+                    ax = axes[i, ch] if num_plots > 1 else axes[ch]
+                    ax.plot(generated_signals[i, ch, :], linewidth=0.8)
+                    if i == 0:
+                        ax.set_title(channel_names[ch], fontweight='bold')
+                    if i == num_plots - 1:
+                        ax.set_xlabel('Time Steps', fontweight='bold')
+                    ax.set_ylabel('Amplitude', fontsize=9)
+                    ax.grid(True, alpha=0.3)
+            
+            fig.suptitle(f'Generated Samples for Station {station_id_str}', fontsize=12, fontweight='bold')
+            plt.tight_layout()
+    
+    # Save to BytesIO buffer and convert to PIL Image, then numpy array
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
+    plt.close(fig)
+    buf.seek(0)
+    
+    # Convert to PIL Image then numpy array (RGB format)
+    img = Image.open(buf)
+    img_array = np.array(img)
+    buf.close()
+    
+    return img_array
+
+# Gradio Interface
+demo = gr.Interface(
+    fn=generate_seismic_data,
+    inputs=[
+        gr.Dropdown(choices=TEST_STATIONS, label="Station ID", value=TEST_STATIONS[0], info="Select one of the 5 test stations"),
+        gr.Slider(minimum=1, maximum=50, value=3, step=1, label="Number of Samples to Generate")
+    ],
+    outputs=gr.Image(label="Generated Seismic Signals", type="numpy"),
+    title="TimesNet-Gen: Site-Specific Strong Motion Generation",
+    description="Generate synthetic seismic signals using Phase 1 model and pre-computed latent bank. Select a station ID and number of samples to generate. (Matching GitHub generate_samples_git.py workflow)"
+)
+
+if __name__ == "__main__":
+    demo.launch()

generate_samples_git.py

@@ -0,0 +1,815 @@
+#!/usr/bin/env python3
+"""
+Simplified inference script for TimesNet-Gen.
+Only loads data for the 5 fine-tuned stations.
+
+Usage:
+    python generate_samples.py --num_samples 50
+"""
+import os
+import argparse
+import torch
+import numpy as np
+from datetime import datetime
+import matplotlib.pyplot as plt
+import glob
+import scipy.io as sio
+
+
+class SimpleArgs:
+    """Configuration for generation."""
+    def __init__(self):
+        # Model architecture
+        self.seq_len = 6000
+        self.d_model = 128
+        self.d_ff = 256
+        self.e_layers = 2
+        self.d_layers = 2
+        self.num_kernels = 6
+        self.top_k = 2
+        self.dropout = 0.1
+        self.latent_dim = 256
+        
+        # System
+        self.use_gpu = torch.cuda.is_available()
+        self.seed = 0
+        
+        # Point-cloud generation
+        self.pcgen_k = 5
+        self.pcgen_jitter_std = 0.0
+
+
+def _iter_np_arrays(obj):
+    """Recursively iterate through numpy arrays in nested structures."""
+    if isinstance(obj, np.ndarray):
+        if obj.dtype == object:
+            for item in obj.flat:
+                yield from _iter_np_arrays(item)
+        else:
+            yield obj
+    elif isinstance(obj, dict):
+        for v in obj.values():
+            yield from _iter_np_arrays(v)
+    elif isinstance(obj, np.void):
+        if obj.dtype.names:
+            for name in obj.dtype.names:
+                yield from _iter_np_arrays(obj[name])
+
+
+def _find_3ch_from_arrays(arrays):
+    """Find 3-channel array from list of arrays."""
+    # Prefer arrays that are 2D with a 3-channel dimension
+    for arr in arrays:
+        if isinstance(arr, np.ndarray) and arr.ndim == 2 and (arr.shape[0] == 3 or arr.shape[1] == 3):
+            return arr
+    # Otherwise, try to find three 1D arrays of same length
+    one_d = [a for a in arrays if isinstance(a, np.ndarray) and a.ndim == 1]
+    for i in range(len(one_d)):
+        for j in range(i + 1, len(one_d)):
+            for k in range(j + 1, len(one_d)):
+                if one_d[i].shape[0] == one_d[j].shape[0] == one_d[k].shape[0]:
+                    return np.stack([one_d[i], one_d[j], one_d[k]], axis=0)
+    return None
+
+
+def load_mat_file(filepath, seq_len=6000, debug=False):
+    """Load and preprocess a .mat file (using data_loader_gen.py logic)."""
+    try:
+        if debug:
+            print(f"\n[DEBUG] Loading: {os.path.basename(filepath)}")
+        
+        # Load with squeeze_me and struct_as_record like data_loader_gen.py
+        mat = sio.loadmat(filepath, squeeze_me=True, struct_as_record=False)
+        
+        if debug:
+            print(f"[DEBUG] Keys in mat file: {[k for k in mat.keys() if not k.startswith('__')]}")
+        
+        # Check if 'EQ' is a struct with nested 'anEQ' structure (like in data_loader_gen.py)
+        if 'EQ' in mat:
+            try:
+                eq_obj = mat['EQ']
+                
+                if debug:
+                    print(f"[DEBUG] EQ type: {type(eq_obj)}")
+                    print(f"[DEBUG] EQ shape: {eq_obj.shape if hasattr(eq_obj, 'shape') else 'N/A'}")
+                
+                # Since struct_as_record=False, EQ is a mat_struct object
+                # Access with attributes, not subscripts
+                if hasattr(eq_obj, 'anEQ'):
+                    dataset = eq_obj.anEQ
+                    if debug:
+                        print(f"[DEBUG] Found anEQ, type: {type(dataset)}")
+                    
+                    if hasattr(dataset, 'Accel'):
+                        accel = dataset.Accel
+                        
+                        if debug:
+                            print(f"[DEBUG] Found Accel: type={type(accel)}, shape={accel.shape if hasattr(accel, 'shape') else 'N/A'}")
+                        
+                        if isinstance(accel, np.ndarray):
+                            # Transpose to (3, N) if needed
+                            if accel.ndim == 2:
+                                if accel.shape[1] == 3:
+                                    accel = accel.T
+                                
+                                if accel.shape[0] == 3:
+                                    data = accel
+                                    if debug:
+                                        print(f"[DEBUG] ✅ Successfully extracted 3-channel data! Shape: {data.shape}")
+                                    
+                                    # Resample if needed
+                                    if data.shape[1] != seq_len:
+                                        from scipy import signal as sp_signal
+                                        data_resampled = np.zeros((3, seq_len), dtype=np.float32)
+                                        for i in range(3):
+                                            data_resampled[i] = sp_signal.resample(data[i], seq_len)
+                                        data = data_resampled
+                                        if debug:
+                                            print(f"[DEBUG] Resampled to {seq_len} samples")
+                                    
+                                    return torch.FloatTensor(data)
+                                else:
+                                    if debug:
+                                        print(f"[DEBUG] Unexpected Accel shape[0]: {accel.shape[0]} (expected 3)")
+                            else:
+                                if debug:
+                                    print(f"[DEBUG] Accel is not 2D: ndim={accel.ndim}")
+                    else:
+                        if debug:
+                            print(f"[DEBUG] anEQ has no 'Accel' attribute")
+                            if hasattr(dataset, '__dict__'):
+                                print(f"[DEBUG] anEQ attributes: {list(vars(dataset).keys())}")
+                else:
+                    if debug:
+                        print(f"[DEBUG] EQ has no 'anEQ' attribute")
+                        if hasattr(eq_obj, '__dict__'):
+                            print(f"[DEBUG] EQ attributes: {list(vars(eq_obj).keys())}")
+                        
+            except Exception as e:
+                if debug:
+                    import traceback
+                    print(f"[DEBUG] Could not parse EQ structure: {e}")
+                    print(f"[DEBUG] Traceback: {traceback.format_exc()}")
+        
+        arrays = list(_iter_np_arrays(mat))
+        
+        if debug:
+            print(f"[DEBUG] Found {len(arrays)} arrays")
+            for i, arr in enumerate(arrays[:5]):  # Show first 5
+                if isinstance(arr, np.ndarray):
+                    print(f"[DEBUG]   Array {i}: shape={arr.shape}, dtype={arr.dtype}")
+        
+        # Common direct keys first
+        for key in ['signal', 'data', 'sig', 'x', 'X', 'signal3c', 'acc', 'NS', 'EW', 'UD']:
+            if key in mat and isinstance(mat[key], np.ndarray):
+                arrays.insert(0, mat[key])
+                if debug:
+                    print(f"[DEBUG] Found key '{key}': shape={mat[key].shape}")
+        
+        # Find 3-channel array
+        data = _find_3ch_from_arrays(arrays)
+        
+        if data is None:
+            if debug:
+                print(f"[DEBUG] Could not find 3-channel array!")
+            return None
+        
+        if debug:
+            print(f"[DEBUG] Found 3-channel data: shape={data.shape}")
+        
+        # Ensure shape is (3, N)
+        if data.shape[0] != 3 and data.shape[1] == 3:
+            data = data.T
+            if debug:
+                print(f"[DEBUG] Transposed to: shape={data.shape}")
+        
+        if data.shape[0] != 3:
+            if debug:
+                print(f"[DEBUG] Wrong number of channels: {data.shape[0]}")
+            return None
+        
+        # Resample to seq_len
+        if data.shape[1] != seq_len:
+            from scipy import signal as sp_signal
+            data_resampled = np.zeros((3, seq_len), dtype=np.float32)
+            for i in range(3):
+                data_resampled[i] = sp_signal.resample(data[i], seq_len)
+            data = data_resampled
+            if debug:
+                print(f"[DEBUG] Resampled to: shape={data.shape}")
+        
+        if debug:
+            print(f"[DEBUG] ✅ Successfully loaded!")
+        
+        return torch.FloatTensor(data)
+    
+    except Exception as e:
+        if debug:
+            print(f"[DEBUG] ❌ Exception: {e}")
+        return None
+
+
+def load_model(checkpoint_path, args):
+    """Load pre-trained TimesNet-PointCloud model."""
+    from models.TimesNet_PointCloud import TimesNetPointCloud
+    
+    # Create model config
+    class ModelConfig:
+        def __init__(self, args):
+            self.seq_len = args.seq_len
+            self.pred_len = 0
+            self.enc_in = 3
+            self.c_out = 3
+            self.d_model = args.d_model
+            self.d_ff = args.d_ff
+            self.num_kernels = args.num_kernels
+            self.top_k = args.top_k
+            self.e_layers = args.e_layers
+            self.d_layers = args.d_layers
+            self.dropout = args.dropout
+            self.embed = 'timeF'
+            self.freq = 'h'
+            self.latent_dim = args.latent_dim
+    
+    config = ModelConfig(args)
+    model = TimesNetPointCloud(config)
+    
+    # Load checkpoint
+    checkpoint = torch.load(checkpoint_path, map_location='cpu')
+    if 'model_state_dict' in checkpoint:
+        model.load_state_dict(checkpoint['model_state_dict'])
+    else:
+        model.load_state_dict(checkpoint)
+    
+    model.eval()
+    if args.use_gpu:
+        model = model.cuda()
+    
+    print(f"[INFO] Model loaded successfully from {checkpoint_path}")
+    return model
+
+
+def generate_samples_from_latent_bank(model, latent_bank_path, station_id, num_samples, args, encoder_std=None):
+    """
+    Generate samples directly from pre-computed latent bank.
+    NO REAL DATA NEEDED!
+    
+    Args:
+        model: TimesNet model
+        latent_bank_path: Path to latent_bank_phase1.npz
+        station_id: Station ID (e.g., '0205')
+        num_samples: Number of samples to generate
+        args: Model arguments
+        encoder_std: Encoder std vector for noise injection
+    
+    Returns:
+        generated_signals: (num_samples, 3, seq_len) array
+        real_names_used: List of lists indicating which latent vectors were used
+    """
+    print(f"[INFO] Loading latent bank from {latent_bank_path}...")
+    
+    try:
+        latent_data = np.load(latent_bank_path)
+    except Exception as e:
+        print(f"[ERROR] Could not load latent bank: {e}")
+        return None, None
+    
+    # Load latent vectors for this station
+    latents_key = f'latents_{station_id}'
+    means_key = f'means_{station_id}'
+    stdev_key = f'stdev_{station_id}'
+    
+    if latents_key not in latent_data:
+        print(f"[ERROR] Station {station_id} not found in latent bank!")
+        print(f"Available stations: {[k.replace('latents_', '') for k in latent_data.keys() if k.startswith('latents_')]}")
+        return None, None
+    
+    latents = latent_data[latents_key]  # (N_samples, seq_len, d_model)
+    means = latent_data[means_key]      # (N_samples, seq_len, d_model)
+    stdevs = latent_data[stdev_key]     # (N_samples, seq_len, d_model)
+    
+    print(f"[INFO] Loaded {len(latents)} latent vectors for station {station_id}")
+    print(f"[INFO] Generating {num_samples} samples via bootstrap aggregation...")
+    
+    generated_signals = []
+    real_names_used = []
+    
+    model.eval()
+    with torch.no_grad():
+        for i in range(num_samples):
+            # Bootstrap: randomly select k latent vectors with replacement
+            k = min(args.pcgen_k, len(latents))
+            selected_indices = np.random.choice(len(latents), size=k, replace=True)
+            
+            # Mix latent features (average)
+            selected_latents = latents[selected_indices]  # (k, seq_len, d_model)
+            selected_means = means[selected_indices]      # (k, seq_len, d_model)
+            selected_stdevs = stdevs[selected_indices]    # (k, seq_len, d_model)
+
+            mixed_features = np.mean(selected_latents, axis=0)  # (seq_len, d_model)
+            mixed_means = np.mean(selected_means, axis=0)       # (seq_len, d_model)
+            mixed_stdevs = np.mean(selected_stdevs, axis=0)     # (seq_len, d_model)
+
+            # NOTE: Do NOT add noise during generation (matching untitled1_gen.py)
+            # untitled1_gen.py only uses noise during TRAINING (Phase 1), not during generation
+            # if encoder_std is not None:
+            #     noise = np.random.randn(*mixed_features.shape) * encoder_std
+            #     mixed_features = mixed_features + noise
+
+            # Convert to torch tensors
+            mixed_features_torch = torch.from_numpy(mixed_features).float().unsqueeze(0)  # (1, seq_len, d_model)
+            means_b = torch.from_numpy(mixed_means).float().unsqueeze(0)   # (1, seq_len, d_model)
+            stdev_b = torch.from_numpy(mixed_stdevs).float().unsqueeze(0)  # (1, seq_len, d_model)
+
+            if args.use_gpu:
+                mixed_features_torch = mixed_features_torch.cuda()
+                means_b = means_b.cuda()
+                stdev_b = stdev_b.cuda()
+
+            # Decode
+            xg = model.project_features_for_reconstruction(mixed_features_torch, means_b, stdev_b)
+            
+            # Store - transpose to (3, 6000)
+            generated_np = xg.squeeze(0).cpu().numpy().T  # (6000, 3) → (3, 6000)
+            generated_signals.append(generated_np)
+            
+            # Track which latent indices were used
+            real_names_used.append([f"latent_{idx}" for idx in selected_indices])
+            
+            if (i + 1) % 10 == 0:
+                print(f"  Generated {i + 1}/{num_samples} samples...")
+    
+    return np.array(generated_signals), real_names_used
+
+
+def _preprocess_component_boore(data: np.ndarray, fs: float, corner_freq: float, filter_order: int = 2) -> np.ndarray:
+    """Boore (2005) style preprocessing: detrend (linear), zero-padding, high-pass Butterworth (zero-phase)."""
+    from scipy.signal import butter, filtfilt
+    x = np.asarray(data, dtype=np.float64)
+    n = x.shape[0]
+    # Linear detrend
+    t = np.arange(n, dtype=np.float64)
+    t_mean = t.mean()
+    x_mean = x.mean()
+    denom = np.sum((t - t_mean) ** 2)
+    slope = 0.0 if denom == 0 else float(np.sum((t - t_mean) * (x - x_mean)) / denom)
+    intercept = float(x_mean - slope * t_mean)
+    x_detr = x - (slope * t + intercept)
+    # Zero-padding
+    Tzpad = (1.5 * filter_order) / max(corner_freq, 1e-6)
+    pad_samples = int(round(Tzpad * fs))
+    x_pad = np.concatenate([np.zeros(pad_samples, dtype=np.float64), x_detr, np.zeros(pad_samples, dtype=np.float64)])
+    # High-pass filter (zero-phase)
+    normalized = corner_freq / (fs / 2.0)
+    normalized = min(max(normalized, 1e-6), 0.999999)
+    b, a = butter(filter_order, normalized, btype='high')
+    y = filtfilt(b, a, x_pad)
+    return y
+
+
+def _konno_ohmachi_smoothing(spectrum: np.ndarray, freq: np.ndarray, b: float = 40.0) -> np.ndarray:
+    """Konno-Ohmachi smoothing as in MATLAB reference (O(n^2))."""
+    f = np.asarray(freq, dtype=np.float64).reshape(-1)
+    s = np.asarray(spectrum, dtype=np.float64).reshape(-1)
+    f = np.where(f == 0.0, 1e-12, f)
+    n = f.shape[0]
+    out = np.zeros_like(s)
+    for i in range(n):
+        w = np.exp(-b * (np.log(f / f[i])) ** 2)
+        w[~np.isfinite(w)] = 0.0
+        denom = np.sum(w)
+        out[i] = 0.0 if denom == 0 else float(np.sum(w * s) / denom)
+    return out
+
+
+def _compute_hvsr_simple(signal: np.ndarray, fs: float = 100.0):
+    """Compute HVSR curve using MATLAB-style pipeline (Boore HP filter + FAS + Konno-Ohmachi)."""
+    try:
+        if signal.ndim != 2 or signal.shape[1] != 3:
+            return None, None
+        if np.any(np.isnan(signal)) or np.any(np.isinf(signal)):
+            return None, None
+
+        # Preprocess components (Boore 2005): detrend + zero-padding + high-pass (0.05 Hz)
+        ew = _preprocess_component_boore(signal[:, 0], fs, 0.05, 2)
+        ns = _preprocess_component_boore(signal[:, 1], fs, 0.05, 2)
+        ud = _preprocess_component_boore(signal[:, 2], fs, 0.05, 2)
+
+        n = int(min(len(ew), len(ns), len(ud)))
+        if n < 16:
+            return None, None
+        ew = ew[:n]; ns = ns[:n]; ud = ud[:n]
+
+        # FFT amplitudes and linear frequency grid
+        half = n // 2
+        if half <= 1:
+            return None, None
+        freq = (np.arange(0, half, dtype=np.float64)) * (fs / n)
+        amp_ew = np.abs(np.fft.fft(ew))[:half]
+        amp_ns = np.abs(np.fft.fft(ns))[:half]
+        amp_ud = np.abs(np.fft.fft(ud))[:half]
+
+        # Horizontal combination via geometric mean, then Konno-Ohmachi smoothing
+        combined_h = np.sqrt(np.maximum(amp_ew, 0.0) * np.maximum(amp_ns, 0.0))
+        sm_h = _konno_ohmachi_smoothing(combined_h, freq, 40.0)
+        sm_v = _konno_ohmachi_smoothing(amp_ud, freq, 40.0)
+
+        sm_v_safe = np.where(sm_v <= 0.0, 1e-12, sm_v)
+        sm_hvsr = sm_h / sm_v_safe
+
+        # Limit to 1-20 Hz band
+        mask = (freq >= 1.0) & (freq <= 20.0)
+        if not np.any(mask):
+            return None, None
+        return freq[mask], sm_hvsr[mask]
+    except Exception:
+        return None, None
+
+
+def save_generated_samples(generated_signals, real_names, station_id, output_dir):
+    """Save generated samples to NPZ file with HVSR and f0 data."""
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Compute HVSR and f0 for all generated signals
+    f0_list = []
+    hvsr_curves = []
+    fs = 100.0
+    
+    print(f"[INFO] Computing HVSR and f0 for {len(generated_signals)} generated samples...")
+    for idx, sig in enumerate(generated_signals):
+        # sig is (3, T), need to transpose to (T, 3)
+        sig_t = sig.T  # (T, 3)
+        freq, hvsr = _compute_hvsr_simple(sig_t, fs)
+        if freq is not None and hvsr is not None:
+            hvsr_curves.append((freq, hvsr))
+            # f0 = frequency at max HVSR
+            max_idx = np.argmax(hvsr)
+            f0 = float(freq[max_idx])
+            f0_list.append(f0)
+    
+    # Build median HVSR curve on a fixed frequency grid (1-20 Hz, 400 points for consistency)
+    hvsr_freq = None
+    hvsr_median = None
+    if hvsr_curves:
+        # Use a fixed frequency grid for consistency with other plots
+        hvsr_freq = np.linspace(1.0, 20.0, 400)
+        # Interpolate all curves to common grid
+        hvsr_matrix = []
+        for freq, hvsr in hvsr_curves:
+            hvsr_interp = np.interp(hvsr_freq, freq, hvsr, left=hvsr[0], right=hvsr[-1])
+            hvsr_matrix.append(hvsr_interp)
+        hvsr_median = np.median(np.vstack(hvsr_matrix), axis=0)
+    
+    # Build f0 histogram (PDF)
+    f0_bins = np.linspace(1.0, 20.0, 21)
+    f0_array = np.array(f0_list)
+    f0_hist, _ = np.histogram(f0_array, bins=f0_bins)
+    f0_pdf = f0_hist.astype(float)
+    f0_sum = f0_pdf.sum()
+    if f0_sum > 0:
+        f0_pdf = f0_pdf / f0_sum
+    
+    # Save timeseries NPZ with HVSR data
+    output_path = os.path.join(output_dir, f'station_{station_id}_generated_timeseries.npz')
+    np.savez_compressed(
+        output_path,
+        generated_signals=generated_signals,
+        signals_generated=generated_signals,  # Alias for compatibility
+        real_names=real_names,
+        station_id=station_id,
+        station=station_id,  # Alias for compatibility
+        f0_timesnet=f0_array,
+        f0_bins=f0_bins,
+        pdf_timesnet=f0_pdf,
+        hvsr_freq_timesnet=hvsr_freq if hvsr_freq is not None else np.array([]),
+        hvsr_median_timesnet=hvsr_median if hvsr_median is not None else np.array([]),
+    )
+    print(f"[INFO] Saved {len(generated_signals)} generated samples to {output_path}")
+    if len(f0_list) > 0:
+        print(f"[INFO]   - f0 samples: {len(f0_list)}, median f0: {np.median(f0_array):.2f} Hz")
+    else:
+        print(f"[INFO]   - No valid HVSR computed")
+
+
+def fine_tune_model(model, all_station_files, args, encoder_std, epochs=10, lr=1e-4):
+    """
+    Fine-tune the model on 5 stations with noise injection.
+    Matches Phase 1 training in untitled1_gen.py exactly.
+    """
+    print("\n" + "="*80)
+    print("Phase 1: Fine-Tuning with Noise Injection")
+    print("="*80)
+    
+    # Prepare data loader
+    all_data = []
+    for station_id, files in all_station_files.items():
+        for fpath in files:
+            data = load_mat_file(fpath, args.seq_len, debug=False)
+            if data is not None:
+                all_data.append(data)
+    
+    if len(all_data) == 0:
+        print("[WARN] No data loaded for fine-tuning!")
+        return model
+    
+    print(f"[INFO] Loaded {len(all_data)} samples for fine-tuning")
+    
+    # Create optimizer (matching untitled1_gen.py Phase 1)
+    batch_size = 32
+    weight_decay = 1e-4
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
+    
+    # AMP scaler (matching untitled1_gen.py)
+    scaler = torch.cuda.amp.GradScaler(enabled=(args.use_gpu))
+    
+    # Gradient clipping (matching untitled1_gen.py)
+    grad_clip = 1.0
+    
+    train_losses_p1 = []
+    
+    for epoch in range(epochs):
+        model.train()
+        total_loss = 0.0
+        total_rec = 0.0
+        num_batches = 0
+        
+        # Shuffle data
+        np.random.shuffle(all_data)
+        
+        for i in range(0, len(all_data), batch_size):
+            batch = all_data[i:i+batch_size]
+            if len(batch) == 0:
+                continue
+            
+            # Stack batch
+            x_list = []
+            for sig in batch:
+                # sig is (3, 6000), transpose to (6000, 3)
+                x_list.append(sig.transpose(0, 1))
+            
+            x = torch.stack(x_list, dim=0)  # (batch, 6000, 3)
+            if args.use_gpu:
+                x = x.cuda()
+            
+            # Zero gradients (matching untitled1_gen.py)
+            optimizer.zero_grad(set_to_none=True)
+            
+            # Forward with AMP and noise injection (matching untitled1_gen.py Phase 1)
+            with torch.cuda.amp.autocast(enabled=(args.use_gpu)):
+                enc_out, means_b, stdev_b = model.encode_features_for_reconstruction(x)
+                
+                # Add noise if encoder_std available (matching untitled1_gen.py line 945-948)
+                if encoder_std is not None:
+                    std_vec = torch.from_numpy(encoder_std).to(enc_out.device).float()
+                    noise = torch.randn_like(enc_out) * std_vec.view(1, 1, -1) * 1.0  # noise_std_scale=1.0
+                    enc_out = enc_out + noise
+                
+                # Decode
+                x_hat = model.project_features_for_reconstruction(enc_out, means_b, stdev_b)
+                
+                # Reconstruction loss (MSE, matching untitled1_gen.py)
+                loss_rec = torch.nn.functional.mse_loss(x_hat, x)
+                loss = loss_rec
+            
+            # Backward with gradient scaling (matching untitled1_gen.py)
+            scaler.scale(loss).backward()
+            
+            # Gradient clipping (matching untitled1_gen.py)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=grad_clip)
+            
+            # Optimizer step with scaler (matching untitled1_gen.py)
+            scaler.step(optimizer)
+            scaler.update()
+            
+            total_loss += float(loss.detach().cpu())
+            total_rec += float(loss_rec.detach().cpu())
+            num_batches += 1
+        
+        # Scheduler step (matching untitled1_gen.py)
+        scheduler.step()
+        
+        avg_loss = total_loss / max(1, num_batches)
+        avg_rec = total_rec / max(1, num_batches)
+        train_losses_p1.append(avg_loss)
+        print(f"[P1] epoch {epoch+1}/{epochs} loss={avg_loss:.4f} (rec={avg_rec:.4f})")
+    
+    print("[INFO] Phase 1 fine-tuning complete!")
+    
+    # Save fine-tuned model (matching untitled1_gen.py Phase 1 checkpoint)
+    checkpoint_dir = './checkpoints'
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    fine_tuned_path = os.path.join(checkpoint_dir, 'timesnet_pointcloud_phase1_finetuned.pth')
+    torch.save({
+        'epoch': epochs,
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'train_losses_phase1': train_losses_p1,
+        'phase': 'phase1'
+    }, fine_tuned_path)
+    print(f"[INFO] ✓ Fine-tuned model saved to: {fine_tuned_path}")
+    
+    return model
+
+
+def plot_sample_preview(generated_signals, station_id, output_dir, num_preview=2):
+    """Create preview plots."""
+    os.makedirs(output_dir, exist_ok=True)
+    
+    for i in range(min(num_preview, len(generated_signals))):
+        fig, axes = plt.subplots(3, 1, figsize=(12, 8))
+        signal = generated_signals[i]
+        
+        channel_names = ['E-W', 'N-S', 'U-D']
+        for ch, (ax, name) in enumerate(zip(axes, channel_names)):
+            ax.plot(signal[ch], linewidth=0.8)
+            ax.set_ylabel(f'{name}\nAmplitude', fontsize=10, fontweight='bold')
+            ax.grid(True, alpha=0.3)
+        
+        axes[-1].set_xlabel('Time Steps', fontsize=10, fontweight='bold')
+        fig.suptitle(f'Generated Sample - Station {station_id}', fontsize=12, fontweight='bold')
+        plt.tight_layout()
+        
+        output_path = os.path.join(output_dir, f'station_{station_id}_preview_{i}.png')
+        plt.savefig(output_path, dpi=150, bbox_inches='tight')
+        plt.close()
+    
+    print(f"[INFO] Saved {min(num_preview, len(generated_signals))} preview plots to {output_dir}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Generate seismic samples (simplified version)')
+    parser.add_argument('--checkpoint', type=str, 
+                        default=r'D:\Baris\codes\Time-Series-Library-main\checkpoints\timesnet_pointcloud_phase1_final.pth',
+                        help='Path to pre-trained model checkpoint')
+    parser.add_argument('--latent_bank', type=str, 
+                        default=r'D:\Baris\codes\Time-Series-Library-main\checkpoints\latent_bank_phase1.npz',
+                        help='Path to latent bank NPZ file')
+    parser.add_argument('--num_samples', type=int, default=50, 
+                        help='Number of samples to generate per station')
+    parser.add_argument('--output_dir', type=str, default='./generated_samples', 
+                        help='Output directory')
+    parser.add_argument('--num_preview', type=int, default=2, 
+                        help='Number of preview plots per station')
+    parser.add_argument('--stations', type=str, nargs='+', default=['0205', '1716', '2020', '3130', '4628'],
+                        help='Target station IDs')
+    parser.add_argument('--data_root', type=str, default=r"D:\Baris\5stats/", 
+                        help='Root path to seismic data (only needed if --fine_tune is used)')
+    parser.add_argument('--fine_tune', action='store_true',
+                        help='Fine-tune the model before generation (use with Phase 0 checkpoint)')
+    parser.add_argument('--fine_tune_epochs', type=int, default=10,
+                        help='Number of fine-tuning epochs')
+    parser.add_argument('--fine_tune_lr', type=float, default=1e-4,
+                        help='Learning rate for fine-tuning')
+    
+    args_cli = parser.parse_args()
+    
+    # Check checkpoint
+    if not os.path.exists(args_cli.checkpoint):
+        print(f"\n{'='*80}")
+        print(f"❌ ERROR: Checkpoint not found!")
+        print(f"{'='*80}")
+        print(f"\nLooking for: {args_cli.checkpoint}")
+        return
+    
+    # Create configuration
+    args = SimpleArgs()
+    
+    print("="*80)
+    print("TimesNet-Gen Sample Generation (Simplified)")
+    print("="*80)
+    print(f"Checkpoint: {args_cli.checkpoint}")
+    print(f"Target stations: {args_cli.stations}")
+    print(f"Samples per station: {args_cli.num_samples}")
+    print(f"Output directory: {args_cli.output_dir}")
+    print("="*80)
+    
+    # Set random seed
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    
+    # Load model
+    model = load_model(args_cli.checkpoint, args)
+    
+    # Try to load encoder_std from Phase 0 (only needed if fine-tuning)
+    encoder_std_path = './pcgen_stats/encoder_feature_std.npy'
+    encoder_std = None
+    if os.path.exists(encoder_std_path):
+        encoder_std = np.load(encoder_std_path)
+        print(f"[INFO] Loaded encoder_std from {encoder_std_path} (shape: {encoder_std.shape})")
+        print(f"[INFO] encoder_std loaded (used only for fine-tuning, NOT for generation)")
+    else:
+        print(f"[INFO] No encoder_std found (not needed for generation, only for fine-tuning)")
+    
+    # Check if latent bank exists
+    if not os.path.exists(args_cli.latent_bank):
+        print(f"\n❌ ERROR: Latent bank not found!")
+        print(f"Looking for: {args_cli.latent_bank}")
+        print(f"\nPlease run untitled1_gen.py first to generate the latent bank.")
+        return
+    
+    print(f"[INFO] Using latent bank: {args_cli.latent_bank}")
+    
+    # Fine-tune if requested (requires real data)
+    if args_cli.fine_tune:
+        print("\n[INFO] Fine-tuning enabled! Loading real data...")
+        
+        all_station_files = {}
+        for station_id in args_cli.stations:
+            # Find all .mat files for this station
+            pattern = os.path.join(args_cli.data_root, f"*{station_id}*.mat")
+            station_files = glob.glob(pattern)
+            
+            if len(station_files) == 0:
+                print(f"[WARN] No files found for station {station_id}")
+            else:
+                print(f"[INFO] Found {len(station_files)} files for station {station_id}")
+                all_station_files[station_id] = station_files
+        
+        if len(all_station_files) == 0:
+            print(f"\n❌ ERROR: No data files found in {args_cli.data_root}")
+            return
+        
+        model = fine_tune_model(model, all_station_files, args, encoder_std, 
+                               epochs=args_cli.fine_tune_epochs, 
+                               lr=args_cli.fine_tune_lr)
+    
+    # Create output directories
+    npz_output_dir = os.path.join(args_cli.output_dir, 'generated_timeseries_npz')
+    plot_output_dir = os.path.join(args_cli.output_dir, 'preview_plots')
+    
+    # Generate samples for each station (from latent bank)
+    print("\n[INFO] Generating samples from latent bank...")
+    for station_id in args_cli.stations:
+        print(f"\n{'='*60}")
+        print(f"Processing Station: {station_id}")
+        print(f"{'='*60}")
+        
+        generated_signals, real_names = generate_samples_from_latent_bank(
+            model, args_cli.latent_bank, station_id, args_cli.num_samples, args, encoder_std
+        )
+        
+        if generated_signals is not None:
+            # Save to NPZ
+            save_generated_samples(generated_signals, real_names, station_id, npz_output_dir)
+            
+            # Create preview plots
+            plot_sample_preview(generated_signals, station_id, plot_output_dir, args_cli.num_preview)
+    
+    print("\n" + "="*80)
+    print("Generation Complete!")
+    print("="*80)
+    print(f"Generated samples saved to: {npz_output_dir}")
+    print(f"Preview plots saved to: {plot_output_dir}")
+    
+    # Debug: Show how many samples were generated per station
+    print("\n[DEBUG] Generated samples per station:")
+    for station_id in args_cli.stations:
+        npz_path = os.path.join(npz_output_dir, f'station_{station_id}_generated_timeseries.npz')
+        if os.path.exists(npz_path):
+            try:
+                data = np.load(npz_path, allow_pickle=True)
+                if 'signals_generated' in data:
+                    n_samples = data['signals_generated'].shape[0]
+                    print(f"  Station {station_id}: {n_samples} samples")
+            except Exception as e:
+                print(f"  Station {station_id}: Error loading NPZ - {e}")
+    print("="*80)
+    
+    # Create HVSR comparison plots (import plot_combined_hvsr_all_sources and call main)
+    print("\n[INFO] Creating HVSR comparison plots (matrices, HVSR curves, f0 distributions)...")
+    print("[INFO] Only plotting TimesNet-Gen vs Real (no Recon/VAE)")
+    try:
+        import sys
+        # Import the plotting module
+        import plot_combined_hvsr_all_sources as hvsr_plotter
+        
+        # Override sys.argv to pass arguments to the plotter
+        # Only provide gen_dir and gen_ts_dir, explicitly disable others with empty strings
+        original_argv = sys.argv
+        sys.argv = [
+            'plot_combined_hvsr_all_sources.py',
+            '--gen_dir', npz_output_dir,  # Use our generated NPZs as gen_dir (they now have HVSR/f0 data)
+            '--gen_ts_dir', npz_output_dir,  # Also use for timeseries plots
+            '--out', os.path.join(args_cli.output_dir, 'hvsr_analysis'),
+            '--recon_dir', '',  # Explicitly empty to disable auto-default
+            '--vae_dir', '',    # Explicitly empty to disable auto-default
+            '--vae_gen_dir', '',  # Explicitly empty to disable auto-default
+        ]
+        
+        # Call the main plotting function
+        hvsr_plotter.main()
+        
+        # Restore original argv
+        sys.argv = original_argv
+        
+        print(f"[INFO] ✅ HVSR analysis complete! Plots saved to: {os.path.join(args_cli.output_dir, 'hvsr_analysis')}")
+    except Exception as e:
+        import traceback
+        print(f"[WARN] Could not create HVSR plots: {e}")
+        traceback.print_exc()
+
+
+if __name__ == "__main__":
+    main()
+