dsipts 1.1.5__py3-none-any.whl

dsipts/models/TimeXER.py

@@ -0,0 +1,184 @@
+## Copyright https://github.com/thuml/Time-Series-Library/blob/main/models/TimeMixer.py
+## Modified for notation alignmenet and batch structure
+## extended to what inside timexer folder
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+
+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 .utils import QuantileLossMO,Permute, get_activation
+from .itransformer.SelfAttention_Family import FullAttention, AttentionLayer
+from .itransformer.Embed import DataEmbedding_inverted
+from .timexer.Layers import FlattenHead,EnEmbedding, EncoderLayer, Encoder
+
+from typing import List, Union
+from ..data_structure.utils import beauty_string
+from .utils import  get_scope
+from .utils import Embedding_cat_variables
+
+
+
+
+class TimeXER(Base):
+    handle_multivariate = True
+    handle_future_covariates = True # or at least it seems...
+    handle_categorical_variables = True #solo nel encoder
+    handle_quantile_loss = True # NOT EFFICIENTLY ADDED, TODO fix this
+    description = get_scope(handle_multivariate,handle_future_covariates,handle_categorical_variables,handle_quantile_loss)
+    
+    def __init__(self, 
+                 patch_len:int,
+                 d_model: int,
+                 n_head: int,
+                 d_ff:int=512,
+                 dropout_rate: float=0.1,
+                 n_layer_decoder: int=1,
+                 activation: str='',
+                 **kwargs)->None:
+        """Initialize the model with specified parameters. https://github.com/thuml/Time-Series-Library/blob/main/models/TimeMixer.py
+        
+        Args:
+            patch_len (int): Length of the patches.
+            d_model (int): Dimension of the model.
+            n_head (int): Number of attention heads.
+            d_ff (int, optional): Dimension of the feedforward network. Defaults to 512.
+            dropout_rate (float, optional): Dropout rate for regularization. Defaults to 0.1.
+            n_layer_decoder (int, optional): Number of layers in the decoder. Defaults to 1.
+            activation (str, optional): Activation function to use. Defaults to ''.
+            **kwargs: Additional keyword arguments passed to the superclass.
+        
+        Raises:
+            ValueError: If an invalid activation function is provided.
+        
+   
+        """
+        super().__init__(**kwargs)
+        if activation == 'torch.nn.SELU':
+            beauty_string('SELU do not require BN','info',self.verbose)
+            use_bn = False
+        if isinstance(activation,str):
+            activation = get_activation(activation)
+        self.save_hyperparameters(logger=False)
+
+
+                
+
+
+        self.patch_len = patch_len
+        self.patch_num = int(self.past_steps // patch_len)
+        d_model = d_model*self.mul
+        
+        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.output_attention = False## not need output attention
+        
+        ###
+        self.en_embedding = EnEmbedding(self.past_channels, d_model, patch_len, dropout_rate)
+
+        self.ex_embedding = DataEmbedding_inverted(self.past_steps, d_model, embed_type='what?', freq='what?', dropout=dropout_rate)  ##embed, freq not used inside
+
+
+        # Encoder-only architecture
+        self.encoder = Encoder(
+            [
+                EncoderLayer(
+                    AttentionLayer(
+                        FullAttention(False, factor = 0.1, attention_dropout=dropout_rate, ##NB factor is not used
+                                      output_attention=False),
+                        d_model, n_head),
+                    AttentionLayer(
+                        FullAttention(False, 0.1, attention_dropout=dropout_rate,
+                                      output_attention=False),
+                        d_model, n_head),
+                    d_model,
+                    d_ff,
+                    dropout=dropout_rate,
+                    activation=activation(),
+                )
+                for l in range(n_layer_decoder)
+            ],
+            norm_layer=torch.nn.LayerNorm(d_model)
+        )
+        self.head_nf = d_model * (self.patch_num + 1)
+        self.head = FlattenHead(self.past_channels, self.head_nf, self.future_steps*self.mul,  head_dropout=dropout_rate)
+        
+        
+        self.future_reshape = nn.Linear(self.future_steps,self.future_steps*self.mul)
+        self.final_linear = nn.Sequential(activation(),
+            nn.Linear(self.past_channels+self.future_channels+emb_fut_out_channel,(self.past_channels+self.future_channels+emb_fut_out_channel)//2),
+            activation(),
+            nn.Linear((self.past_channels+self.future_channels+emb_fut_out_channel)//2,self.out_channels)  
+        )
+
+        
+
+
+
+
+    def forward(self, batch:dict)-> float:
+
+
+        x_enc = batch['x_num_past'].to(self.device)
+
+        BS = x_enc.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)
+        tmp_future = [emb_fut]
+        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)
+
+        if 'x_num_future' in batch.keys():
+            x_future = batch['x_num_future'].to(self.device)
+            tmp_future.append(x_future)
+        if len(tmp_future)>0:
+            tmp_future = torch.cat(tmp_future,2)
+        else:
+            tmp_future = None
+
+
+
+
+
+        en_embed, n_vars = self.en_embedding(x_enc.permute(0, 2, 1))
+        ex_embed = self.ex_embedding(x_enc, emb_past)
+
+        enc_out = self.encoder(en_embed, ex_embed)
+        
+        
+
+
+        enc_out = torch.reshape(enc_out, (-1, n_vars, enc_out.shape[-2], enc_out.shape[-1]))
+        
+
+        # z: [bs x nvars x d_model x patch_num]
+        enc_out = enc_out.permute(0, 1, 3, 2)
+
+        dec_out = self.head(enc_out)  # z: [bs x nvars x target_window]
+        #dec_out = dec_out.permute(0, 2, 1)
+        if tmp_future is not None:
+            tmp_future = self.future_reshape(tmp_future.permute(0, 2, 1))
+            dec_out = torch.cat([tmp_future,dec_out],1)
+        dec_out = self.final_linear(dec_out.permute(0, 2, 1))
+        return dec_out.reshape(BS,self.future_steps,self.out_channels,self.mul)
+        
+        #idx_target = batch['idx_target'][0]
+        #return dec_out[:, :,idx_target].reshape(BS,self.future_steps,self.out_channels,self.mul)
+        
+

dsipts/models/VQVAEA.py

@@ -0,0 +1,299 @@
+
+## Copyright Copyright (c) 2020 Andrej Karpathy https://github.com/karpathy/minGPT?tab=MIT-1-ov-file#readme
+## Modified for notation alignmenet,  batch structure and adapted for timeseries
+## extended to what inside vva folder
+
+from torch import  nn
+import torch
+from torch.nn import functional as F
+
+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 .vva.minigpt import Block
+from .vva.vqvae import VQVAE
+import logging
+from random import random
+from ..data_structure.utils import beauty_string
+from .utils import  get_scope
+
+torch.autograd.set_detect_anomaly(True)
+
+        
+
+class VQVAEA(Base):
+    handle_multivariate = False
+    handle_future_covariates = False
+    handle_categorical_variables = False
+    handle_quantile_loss = False
+    description = get_scope(handle_multivariate,handle_future_covariates,handle_categorical_variables,handle_quantile_loss)
+    
+    
+    def __init__(self, 
+                 past_steps:int,
+                 future_steps:int,
+                 past_channels:int,
+                 future_channels:int,
+                 hidden_channels:int,
+                 embs:List[int],
+                 d_model:int,
+                 max_voc_size:int,
+                 num_layers:int,
+                 dropout_rate:float,
+                 commitment_cost:float,
+                 decay:float,
+                 n_heads:int,
+                 out_channels:int,
+                 epoch_vqvae: int,
+                 persistence_weight:float=0.0,
+                 loss_type: str='l1',
+                 quantiles:List[int]=[],
+                 optim:Union[str,None]=None,
+                 optim_config:dict=None,
+                 scheduler_config:dict=None,
+                 **kwargs)->None:
+        """ Custom encoder-decoder 
+        
+        Args:
+            past_steps (int):  number of past datapoints used 
+            future_steps (int): number of future lag to predict
+            past_channels (int): number of numeric past variables, must be >0
+            future_channels (int): number of future numeric variables 
+            embs (List): list of the initial dimension of the categorical variables
+            cat_emb_dim (int): final dimension of each categorical variable
+            hidden_RNN (int): hidden size of the RNN block
+            num_layers_RNN (int): number of RNN layers
+            kind (str): one among GRU or LSTM
+            kernel_size (int): kernel size in the encoder convolutional block
+            sum_emb (bool): if true the contribution of each embedding will be summed-up otherwise stacked
+            out_channels (int):  number of output channels
+            activation (str, optional): activation fuction function pytorch. Default torch.nn.ReLU
+            remove_last (bool, optional): if True the model learns the difference respect to the last seen point
+            persistence_weight (float):  weight controlling the divergence from persistence model. Default 0
+            loss_type (str, optional): this model uses custom losses or l1 or mse. Custom losses can be linear_penalization or exponential_penalization. Default l1,
+            quantiles (List[int], optional): we can use quantile loss il len(quantiles) = 0 (usually 0.1,0.5, 0.9) or L1loss in case len(quantiles)==0. Defaults to [].
+            dropout_rate (float, optional): dropout rate in Dropout layers
+            use_bn (bool, optional): if true BN layers will be added and dropouts will be removed
+            use_glu (bool,optional): use GLU for feature selection. Defaults to True.
+            glu_percentage (float, optiona): percentage of features to use. Defaults to 1.0.
+            n_classes (int): number of classes (0 in regression)
+            optim (str, optional): if not None it expects a pytorch optim method. Defaults to None that is mapped to Adam.
+            optim_config (dict, optional): configuration for Adam optimizer. Defaults to None.
+            scheduler_config (dict, optional): configuration for stepLR scheduler. Defaults to None.
+
+        """
+
+        
+        super().__init__(**kwargs)
+        self.save_hyperparameters(logger=False)
+        self.d_model = d_model
+        self.max_voc_size = max_voc_size
+        self.future_steps = future_steps
+        self.epoch_vqvae = epoch_vqvae
+        ##PRIMA VQVAE
+        assert out_channels==1, beauty_string('Working only for one singal','section',True)
+        assert past_steps%2==0 and future_steps%2==0, beauty_string('There are some issue with the deconder in case of odd length','section',True)
+        self.vqvae = VQVAE(in_channels=1, hidden_channels=hidden_channels,out_channels=1,num_embeddings= max_voc_size,embedding_dim=d_model,commitment_cost=commitment_cost,decay=decay  )
+        
+        ##POI GPT
+
+
+        self.block_size = past_steps//2 + future_steps//2 -1
+        self.sentence_length = future_steps//2
+        
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(max_voc_size, d_model),
+            wpe = nn.Embedding(self.block_size, d_model),
+            drop = nn.Dropout(dropout_rate),
+            h = nn.ModuleList([Block( d_model,dropout_rate,n_heads,dropout_rate,self.block_size) for _ in range(num_layers)]), ##care can be different dropouts
+            ln_f = nn.LayerNorm(d_model),
+            lm_head = nn.Linear(d_model, max_voc_size, bias=False)
+        ))
+        # report number of parameters (note we don't count the decoder parameters in lm_head)
+        n_params = sum(p.numel() for p in self.transformer.parameters())
+        beauty_string("number of parameters: %.2fM" % (n_params/1e6,),'info',self.verbose)
+        
+        self.use_quantiles = False
+        self.is_classification = True
+        self.optim_config = optim_config
+       
+        
+
+        
+    def configure_optimizers(self):
+        
+    
+        #return torch.optim.Adam(self.vqvae.parameters(), lr=self.optim_config.lr_vqvae,
+        #                                   weight_decay=self.optim_config.weight_decay_vqvae)
+
+
+        return torch.optim.AdamW([
+                {'params':self.vqvae.parameters(),'lr':self.optim_config.lr_vqvae,'weight_decay':self.optim_config.weight_decay_vqvae},
+                {'params':self.transformer.parameters(),'lr':self.optim_config.lr_gpt,'weight_decay':self.optim_config.weight_decay_gpt},
+            ])
+       
+
+        
+    def gpt(self,tokens):
+    
+    
+        b, t = tokens.size()
+        assert t <= self.block_size, beauty_string("Cannot forward sequence of length {t}, block size is only {self.block_size}",'section',True)
+        pos = torch.arange(0, t, dtype=torch.long, device=self.device).unsqueeze(0) # shape (1, t)
+
+        # forward the GPT model itself
+        tok_emb = self.transformer.wte(tokens) # token embeddings of shape (b, t, n_embd)
+        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (1, t, n_embd)
+        x = self.transformer.drop(tok_emb + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+        logits = self.transformer.lm_head(x)
+        return logits
+
+    def forward(self, batch):
+        
+        ##VQVAE
+        #current_epoch = self.current_epoch 
+        #if current_epoch < 1000:
+        #    self.vqvae.train()
+        #    loss_gpt = 100
+        #else:
+        #    self.vqvae.eval()
+        idx_target = batch['idx_target'][0]
+
+        #(tensor([194, 163, 174, 176, 160, 168, 175]),
+        # tensor([ -1,  -1,  -1, 160, 168, 175, 160]))
+
+           
+        data = batch['x_num_past'][:,:,idx_target]
+        if self.current_epoch > self.epoch_vqvae:
+            with torch.no_grad():
+                vqloss, data_recon, perplexity,quantized_x,encodings_x = self.vqvae(data.permute(0,2,1))
+                loss_vqvae = 0
+        else:
+            vq_loss, data_recon, perplexity,quantized_x,encodings_x = self.vqvae(data.permute(0,2,1))
+            if random()<0.001:
+                beauty_string(perplexity,'info',self.verbose)
+            recon_error = F.mse_loss(data_recon.squeeze(), data.squeeze()) 
+            loss_vqvae = recon_error + vq_loss
+
+        if self.current_epoch > self.epoch_vqvae:
+            with torch.no_grad():
+                _, _, _,quantized_y,encodings_y = self.vqvae(batch['y'].permute(0,2,1))
+            
+            ##GPT
+            tokens = torch.cat([encodings_x.argmax(dim=2),encodings_y.argmax(dim=2)[:,0:-1]],1)
+            tokens_y = torch.cat([encodings_x.argmax(dim=2)[:,0:-1],encodings_y.argmax(dim=2)],1)
+            tokens_y[:,0:encodings_x.shape[1]-1] = -1
+            logits = self.gpt(tokens)
+            loss_gpt = F.cross_entropy(logits.view(-1, logits.size(-1)),tokens_y.view(-1), ignore_index=-1)
+
+            
+            ##adesso devo ricostruire la y perche' e quello che voglio come output
+            with torch.no_grad():
+                encoding_indices = torch.argmax(logits.reshape(-1,self.max_voc_size), dim=1).unsqueeze(1) ##
+                encodings = torch.zeros(encoding_indices.shape[0], self.vqvae._vq_vae._num_embeddings, device=self.device)
+                encodings.scatter_(1, encoding_indices, 1)
+                quantized = torch.matmul(encodings, self.vqvae._vq_vae._embedding.weight).view(data.shape[0],-1,self.d_model) ##B x L x hidden
+                quantized = quantized.permute(0, 2, 1).contiguous()
+                y_hat = self.vqvae._decoder(quantized,False).squeeze()[:,-self.future_steps:]
+            
+            l1_loss = nn.L1Loss()(y_hat,batch['y'].squeeze())
+
+            return y_hat, loss_vqvae+loss_gpt+l1_loss
+        else:
+            return None, loss_vqvae
+
+
+    def training_step(self, batch, batch_idx):
+        """
+        pythotrch lightening stuff
+        
+        :meta private:
+        """
+        _, loss = self(batch)
+        return loss
+    
+    def validation_step(self, batch, batch_idx):
+        """
+        pythotrch lightening stuff
+        
+        :meta private:
+        """
+        _, loss = self(batch)
+        return loss
+
+
+    def generate(self, idx, max_new_tokens, temperature=1.0, do_sample=False, top_k=None,num_samples=100):
+        """
+        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
+        the sequence max_new_tokens times, feeding the predictions back into the model each time.
+        Most likely you'll want to make sure to be in model.eval() mode of operation for this.
+        """
+        assert do_sample is False,logging.info('NOT IMPLEMENTED YET')
+        if do_sample:
+            
+            idx = idx.repeat(num_samples,1,1)
+            for _ in range(max_new_tokens):
+                tmp = []
+                for i in range(num_samples):
+                    idx_cond = idx[i,:,:] if idx.size(2) <= self.block_size else idx[i,:, -self.block_size:]
+                    logits = self.gpt(idx_cond)
+                    logits = logits[:, -1, :] / temperature
+                    if top_k is not None:
+                        v, _ = torch.topk(logits, top_k)
+                        logits[logits < v[:, [-1]]] = -float('Inf')
+                    probs = F.softmax(logits, dim=-1)
+                    idx_next = torch.multinomial(probs, num_samples=1, replacement=True)
+                    tmp.append(idx_next)
+                tmp = torch.cat(tmp,dim=1).T.unsqueeze(2)
+                idx = torch.cat((idx, tmp), dim=2)
+            return idx
+        else:
+            for _ in range(max_new_tokens):
+                
+                # if the sequence context is growing too long we must crop it at block_size
+                idx_cond = idx if idx.size(1) <= self.block_size else idx[:, -self.block_size:]
+                # forward the model to get the logits for the index in the sequence
+                logits = self.gpt(idx_cond)
+                # pluck the logits at the final step and scale by desired temperature
+                logits = logits[:, -1, :] / temperature
+                # optionally crop the logits to only the top k options
+                if top_k is not None:
+                    v, _ = torch.topk(logits, top_k)
+                    logits[logits < v[:, [-1]]] = -float('Inf')
+                # apply softmax to convert logits to (normalized) probabilities
+                probs = F.softmax(logits, dim=-1)
+                # either sample from the distribution or take the most likely element
+                _, idx_next = torch.topk(probs, k=1, dim=-1)
+                # append sampled index to the running sequence and continue
+                idx = torch.cat((idx, idx_next), dim=1)
+
+            return idx
+
+    def inference(self, batch:dict)->torch.tensor:
+
+        idx_target = batch['idx_target'][0]
+        data = batch['x_num_past'][:,:,idx_target].to(self.device)
+        vq_loss, data_recon, perplexity,quantized_x,encodings_x = self.vqvae(data.permute(0,2,1))
+        x = encodings_x.argmax(dim=2)
+        inp = x[:, :self.sentence_length]
+        # let the model sample the rest of the sequence
+        cat = self.generate(inp, self.sentence_length, do_sample=False) # non riesco a gestirla qui :-)
+        encoding_indices = cat.flatten().unsqueeze(1) ##
+        encodings = torch.zeros(encoding_indices.shape[0], self.vqvae._vq_vae._num_embeddings, device=self.device)
+        encodings.scatter_(1, encoding_indices, 1)
+        quantized = torch.matmul(encodings, self.vqvae._vq_vae._embedding.weight).view(x.shape[0],-1,self.d_model) ##B x L x hidden
+        quantized = quantized.permute(0, 2, 1).contiguous()
+        y_hat = self.vqvae._decoder(quantized,False).squeeze()[:,-self.future_steps:]
+
+        ## BxLxCx3
+        return y_hat.unsqueeze(2).unsqueeze(3)
+