CFORVAE.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat
import math



class Encoder(nn.Module):
    def __init__(self, temp_dim, f_dim):
        super(Encoder, self).__init__()
        self.temporal = nn.Sequential(
                nn.Linear(temp_dim, temp_dim*2),
                nn.ReLU(),
                nn.Linear(temp_dim*2, temp_dim),
                nn.Dropout(0.1)
            )
        
        self.channel = nn.Sequential(
                nn.Linear(f_dim, f_dim*2),
                nn.ReLU(),
                nn.Linear(f_dim*2, f_dim),
                nn.Dropout(0.1)
            )

    def forward(self, x):
        x = self.temporal(x.permute(0,2,1)).permute(0,2,1)
        x = self.channel(x)
       
        return x


class SpectralConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, modes1, modes2):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.modes1 = modes1  # Time axis (sequence)
        self.modes2 = modes2  # Projected feature axis

        self.scale = 1 / (in_channels * out_channels)
        self.weights = nn.Parameter(
            self.scale * torch.randn(in_channels, out_channels, modes1, modes2, dtype=torch.cfloat)
        )

    def compl_mul2d(self, input, weights):
        # Complex multiplication: (B, in_c, T_ft, F_ft) × (in_c, out_c, T_ft, F_ft) → (B, out_c, T_ft, F_ft)
        return torch.einsum("bctf,coxy->boxy", input, weights)

    def forward(self, x):
        # x: (B, in_c, T, F)
        B, C, T, F = x.shape

        x_ft = torch.fft.rfft2(x, norm="ortho")  # → (B, C, T_ft, F_ft)
        T_ft, F_ft = x_ft.shape[-2], x_ft.shape[-1]

        # Clamp to available modes
        modes1 = min(self.modes1, T_ft)
        modes2 = min(self.modes2, F_ft)

        out_ft = torch.zeros(B, self.out_channels, T_ft, F_ft, dtype=torch.cfloat, device=x.device)

        out_ft[:, :, :modes1, :modes2] = self.compl_mul2d(
            x_ft[:, :, :modes1, :modes2],
            self.weights[:, :, :modes1, :modes2]
        )

        x_out = torch.fft.irfft2(out_ft, s=(T, F), norm="ortho")  # → (B, out_c, T, F)
        # print("x:", x.shape)
        # print("x_out: ", x_out.shape)
        return x_out


class FeatureCorrelationFNO2D(nn.Module):
    def __init__(self, in_features, hidden_dim=32, modes1=4, modes2=4):
        super().__init__()
        self.input_proj = nn.Linear(in_features, hidden_dim)
        
        # For time series data, we typically have many time steps but few features
        # Adjust modes2 to be reasonable for the expected input dimensions
        self.fno = SpectralConv2d(hidden_dim, 1, modes1, modes2)
        self.output_proj = nn.Linear(hidden_dim, in_features)
        self.activation = nn.GELU()

    def forward(self, x):
        # x: (B, S, D)
        B, S, D = x.shape
        x = self.input_proj(x)  # (B, S, H)

        x = x.unsqueeze(1)     # (B, 1, S, H)        
        x = self.fno(x)       
        x = x.squeeze(1)        # (B, S, H)
        x = self.activation(x)
        return self.output_proj(x)  # (B, S, D)


class TemporalAttentionPooling(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.query = nn.Parameter(torch.randn(1, d_model))  # learnable query
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, x):  # [B, T, D]
        x_proj = self.proj(x)  # optional projection: [B, T, D]
        q = self.query.expand(x.size(0), -1)  # [B, D]
        attn = torch.bmm(x_proj, q.unsqueeze(-1)).squeeze(-1)  # [B, T]
        attn = torch.softmax(attn, dim=1)
        pooled = torch.bmm(attn.unsqueeze(1), x).squeeze(1)  # [B, D]
        return pooled
    
     
class CFORVAE(nn.Module):
    def __init__(
        self,
        seq_len,
        input_dim,
        pred_len,
        latent_dim=128,
        use_decomposition=False,
        use_feature_correlation=False,
        use_noise = False
    ):
        super().__init__()
        self.seq_len = seq_len
        self.latent_dim = latent_dim
        self.input_dim = input_dim
        self.use_decomposition = use_decomposition
        self.use_feature_correlation = use_feature_correlation
        self.use_noise = use_noise
        self.pred_len = pred_len
        
        if pred_len is not None:
            self.feature_dim_output = 1
            self.projection_seq = nn.Linear(self.seq_len, self.pred_len)   


        self.input_proj = nn.Linear(input_dim, latent_dim)
        
        decoder_input_dim = latent_dim * 2

        if self.use_decomposition:
            self.encoder_season = Encoder(temp_dim=seq_len, f_dim=latent_dim)
            self.encoder_trend = Encoder(temp_dim=seq_len, f_dim=latent_dim)
            
            self.mu_proj_season = nn.Linear(latent_dim, latent_dim)
            self.logvar_proj_season = nn.Linear(latent_dim, latent_dim)
            
            self.mu_proj_trend = nn.Linear(latent_dim, latent_dim)
            self.logvar_proj_trend = nn.Linear(latent_dim, latent_dim)  
            
            self.temporal_pool_trend = TemporalAttentionPooling(latent_dim)
            self.temporal_pool_season = TemporalAttentionPooling(latent_dim)
             
            
        else:
            self.encoder = Encoder(temp_dim=seq_len, f_dim=latent_dim)
             # Latent space
            self.mu_proj = nn.Linear(latent_dim, latent_dim)
            self.logvar_proj = nn.Linear(latent_dim, latent_dim)
            
        # Optional feature correlation module
        if self.use_feature_correlation:
            self.correlation = FeatureCorrelationFNO2D(in_features=latent_dim, modes1=seq_len//2)
        
       
        # # Decoder
        self.decoder = nn.Sequential(
            nn.Linear(decoder_input_dim, latent_dim * 2),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(latent_dim * 2, latent_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(latent_dim, input_dim)
        )
        
        print(f"CFORVAE initialized with flags:")
        print(f"  - Use decomposition: {self.use_decomposition}")
        print(f"  - Use feature_correlation: {self.use_feature_correlation}")
        
  
    def seasonal_trend_decompose(self, x, kernel_size=7):
        # Simple moving average as trend, seasonal = residual
        trend = F.avg_pool1d(x.permute(0, 2, 1), kernel_size, stride=1, padding=kernel_size//2)
        trend = trend.permute(0, 2, 1)
        seasonal = x - trend
        return seasonal, trend
    
    
    def forward(self, x, mask=None):
        B, T, N = x.shape
        
        x_proj = self.input_proj(x) 
        
                
        if self.use_decomposition:
            # Decompose into trend and seasonal
            seasonal, trend = self.seasonal_trend_decompose(x_proj)
            
            if self.use_feature_correlation:
                trend  = trend + self.correlation(trend)

             # Encode both separately
            enc_trend = self.encoder_trend(trend)                
            enc_seasonal = self.encoder_season(seasonal)
            latent = enc_seasonal + enc_trend
            
            mu_season = self.mu_proj_season(self.temporal_pool_season(enc_seasonal))
            logvar_season = self.logvar_proj_season(self.temporal_pool_season(enc_seasonal))   

            mu_trend = self.mu_proj_trend(self.temporal_pool_trend(enc_trend))
            logvar_trend = self.logvar_proj_trend(self.temporal_pool_trend(enc_trend))
            
            # mu_season = self.mu_proj_season(enc_seasonal.mean(dim=1))
            # logvar_season = self.logvar_proj_season(enc_seasonal.mean(dim=1))   
            
            # mu_trend = self.mu_proj_trend(enc_trend.mean(dim=1))
            # logvar_trend = self.logvar_proj_trend(enc_trend.mean(dim=1))
            
            z_seasonal = self.reparameterize(mu_season, logvar_season)
            z_trend = self.reparameterize(mu_trend, logvar_trend)
            
            z_seasonal = z_seasonal.unsqueeze(1).repeat(1, T, 1)
            z_trend = z_trend.unsqueeze(1).repeat(1, T, 1)
            
            z_expanded = z_seasonal + z_trend   
            
            mu = (mu_season + mu_trend) / 2
            logvar = (logvar_season + logvar_trend) / 2        
            
        else:
            latent = self.encoder(x_proj)

            mu = self.mu_proj(latent.mean(dim=1))
            logvar = self.logvar_proj(latent.mean(dim=1))
            z = self.reparameterize(mu, logvar)
            z_expanded = z.unsqueeze(1).repeat(1, T, 1)

        if self.pred_len is not None:
            x_proj = latent
        elif self.training and self.use_decomposition and self.use_feature_correlation and self.use_noise:
            x_proj = x_proj + torch.randn_like(x_proj) * 0.5

        dec_input = torch.cat([z_expanded, x_proj], dim=-1)
        recon_x = self.decoder(dec_input)

        if self.pred_len is not None:
            recon_x = self.projection_seq(recon_x.permute(0, 2, 1)).permute(0, 2, 1)
            
        return recon_x, mu, logvar
    
    def reparameterize(self, mu, logvar):
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            return mu + eps * std
        else:
            return mu

    def get_config(self):
        """Return current configuration for logging/comparison"""
        return {
            'use_decomposition': self.use_decomposition,
            'use_feature_correlation': self.use_feature_correlation,
            'latent_dim': self.latent_dim,
        }

# Example usage with ablation study
def create_model_configs(seq_len, input_dim, pred_len):
    """Create different model configurations for ablation study"""
    base_config = {
        'seq_len': seq_len,
        'pred_len': pred_len,
        'input_dim': input_dim,
        'latent_dim': 128,
    }
    
    # Different configurations to test
    configs = {
         'full_mode': {
            **base_config,
            'use_decomposition': True,
            'use_feature_correlation': True,
        },
    }
    
    return configs

def create_model(config_name, seq_len, input_dim, pred_len=None, use_noise=False):
    """Create model with specified configuration"""
    configs = create_model_configs(seq_len, input_dim, pred_len)
    config = configs[config_name]
    
    model = CFORVAE(**config, use_noise=use_noise)
    return model, config_name