dsipts 1.1.5__py3-none-any.whl

dsipts/models/CrossFormer.py

@@ -0,0 +1,152 @@
+## Copyright 2023 Yunhao Zhang and Junchi Yan (https://github.com/Thinklab-SJTU/Crossformer?tab=Apache-2.0-1-ov-file#readme)
+## Code modified for align the notation and the batch generation
+## extended to all present in crossformer folder
+
+
+
+from torch import  nn
+import torch
+
+try:
+    import lightning.pytorch as pl
+    from .base_v2 import Base
+    OLD_PL = False
+except:
+    import pytorch_lightning as pl
+    OLD_PL = True
+    from .base import Base
+from typing import List,Union
+from einops import  repeat
+from ..data_structure.utils import beauty_string
+from .utils import  get_scope
+from .crossformer.cross_encoder import Encoder
+from .crossformer.cross_decoder import Decoder
+from .crossformer.cross_embed import DSW_embedding
+from .utils import Embedding_cat_variables
+from math import ceil
+from .utils import Embedding_cat_variables
+from .utils import get_activation
+
+  
+
+  
+class CrossFormer(Base):
+    handle_multivariate = True
+    handle_future_covariates = True
+    handle_categorical_variables = True
+    handle_quantile_loss = True
+
+    description = get_scope(handle_multivariate,handle_future_covariates,handle_categorical_variables,handle_quantile_loss)
+    
+    def __init__(self, 
+                 
+                 d_model:int,
+                 hidden_size:int,
+                 n_head:int,
+                 seg_len:int,
+                 n_layer_encoder:int,
+                 win_size:int,
+                 factor:int=10,
+                 dropout_rate:float=0.1,
+                 activation:str='torch.nn.ReLU',
+
+                 **kwargs)->None:
+        """  CrossFormer (https://openreview.net/forum?id=vSVLM2j9eie)
+
+        
+        Args:
+            d_model (int): The dimensionality of the model.
+            hidden_size (int): The size of the hidden layers.
+            n_head (int): The number of attention heads.
+            seg_len (int): The length of the segments.
+            n_layer_encoder (int): The number of layers in the encoder.
+            win_size (int): The size of the window for attention.
+            factor (int, optional): see .crossformer.attn.TwoStageAttentionLayer. Defaults to 10.
+            dropout_rate (float, optional): The dropout rate. Defaults to 0.1.
+            activation (str, optional): The activation function to use. Defaults to 'torch.nn.ReLU'.
+            **kwargs: Additional keyword arguments for the parent class.
+        
+        Returns:
+            None: This method does not return a value.
+        
+        Raises:
+            ValueError: If the activation function is not recognized.
+            
+
+        """
+      
+        if isinstance(activation, str):
+            activation = get_activation(activation)
+        super().__init__(**kwargs)
+        self.save_hyperparameters(logger=False)
+        
+
+        # The padding operation to handle invisible sgemnet length
+        self.pad_past_steps = ceil(1.0 *self.past_steps / seg_len) * seg_len
+        self.pad_future_steps = ceil(1.0 * self.future_steps / seg_len) * seg_len
+        self.past_steps_add = self.pad_past_steps - self.past_steps
+
+        # Embedding
+        self.emb_past = Embedding_cat_variables(self.past_steps,self.emb_dim,self.embs_past, reduction_mode=self.reduction_mode,use_classical_positional_encoder=self.use_classical_positional_encoder,device = self.device)
+        self.emb_fut = Embedding_cat_variables(self.future_steps,self.emb_dim,self.embs_fut, reduction_mode=self.reduction_mode,use_classical_positional_encoder=self.use_classical_positional_encoder,device = self.device)
+        emb_past_out_channel = self.emb_past.output_channels
+        emb_fut_out_channel = self.emb_fut.output_channels
+        self.enc_value_embedding = DSW_embedding(seg_len, d_model)
+        self.enc_pos_embedding = nn.Parameter(torch.randn(1, self.past_channels+emb_past_out_channel, (self.pad_past_steps // seg_len), d_model))
+        self.pre_norm = nn.LayerNorm(d_model)
+
+        ## Custom embeddings ##these are not used in crossformer
+
+
+        self.encoder = Encoder(n_layer_encoder, win_size, d_model, n_head, hidden_size, block_depth = 1, \
+                                    dropout = dropout_rate,in_seg_num = (self.pad_past_steps // seg_len), factor = factor)
+        
+        # Decoder
+        self.dec_pos_embedding = nn.Parameter(torch.randn(1, self.past_channels+emb_past_out_channel, (self.pad_future_steps // seg_len), d_model))
+        self.decoder = Decoder(seg_len, n_layer_encoder + 1, d_model, n_head, hidden_size, dropout_rate, \
+                                    out_seg_num = (self.pad_future_steps // seg_len), factor = factor)
+     
+        dim = self.past_channels+emb_past_out_channel+emb_fut_out_channel+self.future_channels
+        self.final_layer = nn.Sequential(activation(),
+                                         nn.Linear(dim, dim//2),
+                                         activation(),
+                                         nn.Linear(dim//2, self.mul*self.out_channels ))
+        
+                                         
+        
+     
+        
+    def forward(self, batch):
+
+        x_seq = batch['x_num_past'].to(self.device)#[:,:,idx_target]
+        BS = x_seq.shape[0]
+        
+        if 'x_cat_future' in batch.keys():
+            emb_fut = self.emb_fut(BS,batch['x_cat_future'].to(self.device))
+        else:
+            emb_fut = self.emb_fut(BS,None)
+        if 'x_cat_past' in batch.keys():
+            emb_past = self.emb_past(BS,batch['x_cat_past'].to(self.device))
+        else:
+            emb_past = self.emb_past(BS,None)
+            
+        tmp_future = [emb_fut]
+        
+        if 'x_num_future' in batch.keys():
+            x_future = batch['x_num_future'].to(self.device)
+            tmp_future.append(x_future)
+
+        x_seq = torch.cat([x_seq,emb_past],2) 
+        batch_size = x_seq.shape[0]
+        if (self.past_steps_add != 0):
+            x_seq = torch.cat((x_seq[:, :1, :].expand(-1, self.past_steps_add, -1), x_seq), dim = 1)
+        x_seq = self.enc_value_embedding(x_seq)
+        x_seq += self.enc_pos_embedding
+        x_seq = self.pre_norm(x_seq)        
+        enc_out = self.encoder(x_seq)
+        dec_in = repeat(self.dec_pos_embedding, 'b ts_d l d -> (repeat b) ts_d l d', repeat = batch_size)
+        predict_y = self.decoder(dec_in, enc_out)
+        res = predict_y[:, :self.future_steps,:]
+        tmp_future.append(res)
+        res = self.final_layer(torch.cat(tmp_future,2))
+        return res.reshape(BS, -1, self.out_channels,self.mul)

dsipts/models/D3VAE.py

@@ -0,0 +1,196 @@
+## Copiright https://github.com/PaddlePaddle/PaddleSpatial
+## Modified for notation alignmenet,  batch structure 
+## extended to what inside d3vae folder
+
+from torch import  nn,optim
+import torch
+
+try:
+    import lightning.pytorch as pl
+    from .base_v2 import Base
+    OLD_PL = False
+except:
+    import pytorch_lightning as pl
+    OLD_PL = True
+    from .base import Base
+from typing import Union
+from .d3vae.model import diffusion_generate, denoise_net,pred_net
+from .utils import Embedding_cat_variables
+
+
+from torch.optim.lr_scheduler import StepLR
+
+def copy_parameters(
+    net_source: torch.nn.Module,
+    net_dest: torch.nn.Module,
+    strict= True,
+) -> None:
+    """
+    Copies parameters from one network to another.
+    Parameters
+    ----------
+    net_source
+        Input network.
+    net_dest
+        Output network.
+    strict:
+        whether to strictly enforce that the keys
+        in :attr:`state_dict` match the keys returned by this module's
+        :meth:`~torch.nn.Module.state_dict` function. Default: ``True``
+    """
+
+    net_dest.load_state_dict(net_source.state_dict(), strict=strict)
+
+
+class D3VAE(Base):              
+    
+    def __init__(self,
+                 
+                 
+                 scale=0.1,
+                 hidden_size=64,
+                 num_layers=2,
+                 dropout_rate=0.1,
+                 diff_steps=200,
+                 loss_type='kl',
+                 beta_end=0.01,
+                 beta_schedule='linear',
+                 channel_mult = 2,
+                 mult=1,
+                 num_preprocess_blocks=1,
+                 num_preprocess_cells=3,
+                 num_channels_enc=16,
+                 arch_instance = 'res_mbconv',
+                 num_latent_per_group=6,
+                 num_channels_dec=16,
+                 groups_per_scale=2,
+                 num_postprocess_blocks=1,
+                 num_postprocess_cells=2,
+                 beta_start=0,
+                 
+                 freq='h',
+                
+                 **kwargs
+                 )->None:
+        super().__init__(**kwargs)
+        input_dim = self.past_channels
+        sequence_length = self.past_steps
+        prediction_length = self.future_steps
+        target_dim = self.out_channels
+        embedding_dimension = self.emb_dim
+        ##pytotch lightening stuff
+        self.save_hyperparameters(logger=False)
+        
+        
+        
+        
+        self.gen_net = diffusion_generate(target_dim,embedding_dimension,prediction_length,sequence_length,scale,hidden_size,num_layers,dropout_rate,diff_steps,loss_type,beta_end,beta_schedule, channel_mult,mult,
+                 num_preprocess_blocks,num_preprocess_cells,num_channels_enc,arch_instance,num_latent_per_group,num_channels_dec,groups_per_scale,num_postprocess_blocks,num_postprocess_cells).to(self.device)
+        
+        self.denoise_net = denoise_net(target_dim,embedding_dimension,prediction_length,sequence_length,scale,hidden_size,num_layers,dropout_rate,diff_steps,loss_type,beta_end,beta_schedule, channel_mult,mult,
+                 num_preprocess_blocks,num_preprocess_cells,num_channels_enc,arch_instance,num_latent_per_group,num_channels_dec,groups_per_scale,num_postprocess_blocks,num_postprocess_cells,beta_start,input_dim,freq,self.embs_past).to(self.device)
+        self.diff_step = diff_steps
+        self.pred_net = pred_net(target_dim,embedding_dimension,prediction_length,sequence_length,scale,hidden_size,num_layers,dropout_rate,diff_steps,loss_type,beta_end,beta_schedule, channel_mult,mult,
+                 num_preprocess_blocks,num_preprocess_cells,num_channels_enc,arch_instance,num_latent_per_group,num_channels_dec,groups_per_scale,num_postprocess_blocks,num_postprocess_cells,beta_start,input_dim,freq,self.embs_fut).to(self.device)
+        #self.embedding = DataEmbedding(input_dim, embedding_dimension, freq,dropout_rate)
+        
+        self.psi = 0.5
+        self.gamma = 0.01
+        self.lambda1 = 1.0
+
+        
+    def configure_optimizers(self):
+        """
+        Each model has optim_config and scheduler_config
+        
+        :meta private:
+        """
+        optimizer = optim.Adam(self.denoise_net.parameters(),  **self.optim_config)
+        self.lr = self.optim_config['lr']
+        if self.scheduler_config is not None:
+            scheduler = StepLR(optimizer,**self.scheduler_config)
+            return [optimizer], [scheduler]
+        else:
+            return optimizer
+        
+        
+        
+    def training_step(self, batch, batch_idx):
+        """
+        pythotrch lightening stuff
+        
+        :meta private:
+        """
+        sample, y_noisy,recon,loss2,total_c = self(batch)
+        ##self.compute_loss(batch,y_hat)
+        mse_loss = self.loss(sample, y_noisy)
+        loss1 = - torch.mean(torch.sum(recon, dim=[1, 2, 3]))
+        loss = loss1*self.psi + loss2*self.lambda1 + mse_loss - self.gamma*total_c
+        
+        
+        return loss
+        
+    def validation_step(self, batch, batch_idx):
+        """
+        pythotrch lightening stuff
+        
+        :meta private:
+        """
+        
+        copy_parameters(self.denoise_net, self.pred_net)
+        batch_x = batch['x_num_past'].to(self.device)
+        batch_x_mark = batch['x_cat_past'].to(self.device)
+        batch_y = batch['y'].to(self.device)
+
+        import pdb
+        pdb.set_trace()
+        _, out, _, _ = self.pred_net(batch_x, batch_x_mark)
+        mse = self.loss(out.squeeze(1), batch_y)
+        
+        
+        return mse
+
+        
+        
+    def forward(self,batch:dict)->torch.tensor:
+        
+        B = batch['x_num_past'].shape[0]
+
+        
+        t = torch.randint(0, self.diff_step, (B,)).long().to(self.device)
+        
+        batch_x = batch['x_num_past'].to(self.device)
+        x_mark = batch['x_cat_past'].to(self.device)
+        batch_y = batch['y'].to(self.device)
+        
+        output, y_noisy, total_c, _, loss2 = self.denoise_net(batch_x, x_mark, batch_y, t)
+        recon = output.log_prob(y_noisy)
+        sample = output.sample()
+        
+        
+        return sample,y_noisy,recon,loss2,total_c
+
+   
+   
+   
+    
+
+    
+    def inference(self, batch:dict)->torch.tensor:
+        """Care here, we need to implement it because for predicting the N-step it will use the prediction at step N-1. TODO fix if because I did not implement the
+        know continuous variable presence here
+
+        Args:
+            batch (dict): batch of the dataloader
+
+        Returns:
+            torch.tensor: result
+        """
+        copy_parameters(self.denoise_net, self.pred_net)
+
+        batch_x = batch['x_num_past'].float().to(self.device)
+        batch_x_mark = batch['x_cat_past'].to(self.device)
+        _, out, _, _ = self.pred_net(batch_x, batch_x_mark)
+        #import pdb
+        #pdb.set_trace()
+        return torch.permute(out, (0,2,1,3))