dsipts 1.1.5__py3-none-any.whl

dsipts/models/samformer/utils.py

@@ -0,0 +1,154 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from torch.optim import Optimizer
+
+
+def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
+    """
+    A copy-paste from https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
+    """
+    L, S = query.size(-2), key.size(-2)
+    scale_factor = 1 / np.sqrt(query.size(-1)) if scale is None else scale
+    attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
+    if is_causal:
+        assert attn_mask is None
+        temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril(diagonal=0)
+        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+        attn_bias.to(query.dtype)
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
+        else:
+            attn_bias += attn_mask
+    attn_weight = query @ key.transpose(-2, -1) * scale_factor
+    attn_weight += attn_bias
+    attn_weight = torch.softmax(attn_weight, dim=-1)
+    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
+    return attn_weight @ value
+
+
+
+
+class RevIN(nn.Module):
+    """
+    Reversible Instance Normalization (RevIN) https://openreview.net/pdf?id=cGDAkQo1C0p
+    https://github.com/ts-kim/RevIN
+    """
+    def __init__(self, num_features: int, eps=1e-5, affine=True):
+        """
+        :param num_features: the number of features or channels
+        :param eps: a value added for numerical stability
+        :param affine: if True, RevIN has learnable affine parameters
+        """
+        super(RevIN, self).__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.affine = affine
+        if self.affine:
+            self._init_params()
+
+    def forward(self, x, mode:str):
+        if mode == 'norm':
+            self._get_statistics(x)
+            x = self._normalize(x)
+        elif mode == 'denorm':
+            x = self._denormalize(x)
+        else: raise NotImplementedError
+        return x
+
+    def _init_params(self):
+        # initialize RevIN params: (C,)
+        self.affine_weight = nn.Parameter(torch.ones(self.num_features))
+        self.affine_bias = nn.Parameter(torch.zeros(self.num_features))
+
+    def _get_statistics(self, x):
+        dim2reduce = tuple(range(1, x.ndim-1))
+        self.mean = torch.mean(x, dim=dim2reduce, keepdim=True).detach()
+        self.stdev = torch.sqrt(torch.var(x, dim=dim2reduce, keepdim=True, unbiased=False) + self.eps).detach()
+
+    def _normalize(self, x):
+        x = x - self.mean
+        x = x / self.stdev
+        if self.affine:
+            x = x * self.affine_weight
+            x = x + self.affine_bias
+        return x
+
+    def _denormalize(self, x):
+        if self.affine:
+            x = x - self.affine_bias
+            x = x / (self.affine_weight + self.eps*self.eps)
+        x = x * self.stdev
+        x = x + self.mean
+        return x
+    
+
+class SAM(torch.optim.Optimizer):
+    def __init__(self, params, base_optimizer, rho=0.05, adaptive=False, **kwargs):
+        assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"
+
+        defaults = dict(rho=rho, adaptive=adaptive, **kwargs)
+        super(SAM, self).__init__(params, defaults)
+
+        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
+        self.param_groups = self.base_optimizer.param_groups
+        self.defaults.update(self.base_optimizer.defaults)
+
+    @torch.no_grad()
+    def first_step(self, zero_grad=False):
+        grad_norm = self._grad_norm()
+        for group in self.param_groups:
+            scale = group["rho"] / (grad_norm + 1e-12)
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                self.state[p]["old_p"] = p.data.clone()
+                e_w = (torch.pow(p, 2) if group["adaptive"] else 1.0) * p.grad * scale.to(p)
+                p.add_(e_w)  # Perturb weights in the gradient direction
+
+        if zero_grad:
+            self.zero_grad()
+
+    @torch.no_grad()
+    def second_step(self, zero_grad=False):
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                p.data = self.state[p]["old_p"]  # Restore original weights
+
+        self.base_optimizer.step()  # Apply the sharpness-aware update
+
+        if zero_grad:
+            self.zero_grad()
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        assert closure is not None, "Sharpness Aware Minimization requires closure, but it was not provided"
+
+        with torch.enable_grad():
+            closure()  # First forward-backward pass
+
+        self.first_step(zero_grad=True)
+
+        with torch.enable_grad():
+            closure()  # Second forward-backward pass
+
+        self.second_step()
+
+    def _grad_norm(self):
+        shared_device = self.param_groups[0]["params"][0].device
+        grads = [
+            ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
+            for group in self.param_groups for p in group["params"]
+            if p.grad is not None
+        ]
+        return torch.norm(torch.stack(grads), p=2) if grads else torch.tensor(0.0, device=shared_device)
+
+    def load_state_dict(self, state_dict):
+        super().load_state_dict(state_dict)
+        if hasattr(self, "base_optimizer"):  # Ensure base optimizer exists
+            self.base_optimizer.load_state_dict(state_dict["base_optimizer"])

dsipts/models/tft/sub_nn.py

@@ -0,0 +1,234 @@
+import torch
+import torch.nn as nn
+from typing import Union
+
+class embedding_cat_variables(nn.Module):
+    # at the moment cat_past and cat_fut together
+    def __init__(self, seq_len: int, lag: int, d_model: int, emb_dims: list, device):
+        """Class for embedding categorical variables, adding 3 positional variables during forward
+
+        Args:
+            seq_len (int): length of the sequence (sum of past and future steps)
+            lag (int): number of future step to be predicted
+            hiden_size (int): dimension of all variables after they are embedded
+            emb_dims (list): size of the dictionary for embedding. One dimension for each categorical variable
+            device : -
+        """
+        super().__init__()
+        self.seq_len = seq_len
+        self.lag = lag
+        self.device = device
+        self.cat_embeds = emb_dims + [seq_len, lag+1, 2] # 
+        self.cat_n_embd = nn.ModuleList([
+            nn.Embedding(emb_dim, d_model) for emb_dim in self.cat_embeds
+        ])
+
+    def forward(self, x: Union[torch.Tensor,int],device:torch.device) -> torch.Tensor:
+        """All components of x are concatenated with 3 new variables for data augmentation, in the order:
+        - pos_seq: assign at each step its time-position
+        - pos_fut: assign at each step its future position. 0 if it is a past step
+        - is_fut: explicit for each step if it is a future(1) or past one(0)
+
+        Args:
+            x (torch.Tensor): [bs, seq_len, num_vars]
+
+        Returns:
+            torch.Tensor: [bs, seq_len, num_vars+3, n_embd] 
+        """
+        if isinstance(x, int):
+            no_emb = True
+            B = x
+        else:
+            no_emb = False
+            B, _, _ = x.shape
+        
+        pos_seq = self.get_pos_seq(bs=B).to(device)
+        pos_fut = self.get_pos_fut(bs=B).to(device)
+        is_fut = self.get_is_fut(bs=B).to(device)
+
+        if no_emb:
+            cat_vars = torch.cat((pos_seq, pos_fut, is_fut), dim=2)
+        else:
+            cat_vars = torch.cat((x, pos_seq, pos_fut, is_fut), dim=2)
+
+        cat_n_embd = self.get_cat_n_embd(cat_vars)
+        return cat_n_embd
+
+    def get_pos_seq(self, bs):
+        pos_seq = torch.arange(0, self.seq_len)
+        pos_seq = pos_seq.repeat(bs,1).unsqueeze(2).to(self.device)
+        return pos_seq
+    
+    def get_pos_fut(self, bs):
+        pos_fut = torch.cat((torch.zeros((self.seq_len-self.lag), dtype=torch.long),torch.arange(1,self.lag+1)))
+        pos_fut = pos_fut.repeat(bs,1).unsqueeze(2).to(self.device)
+        return pos_fut
+    
+    def get_is_fut(self, bs):
+        is_fut = torch.cat((torch.zeros((self.seq_len-self.lag), dtype=torch.long),torch.ones((self.lag), dtype=torch.long)))
+        is_fut = is_fut.repeat(bs,1).unsqueeze(2).to(self.device)
+        return is_fut
+    
+    def get_cat_n_embd(self, cat_vars):
+        cat_n_embd = torch.Tensor().to(cat_vars.device)
+        for index, layer in enumerate(self.cat_n_embd):
+            emb = layer(cat_vars[:, :, index])
+            cat_n_embd = torch.cat((cat_n_embd, emb.unsqueeze(2)),dim=2)
+        return cat_n_embd
+
+class LSTM_Model(nn.Module):
+    def __init__(self, num_var: int, d_model: int, pred_step: int, num_layers: int, dropout: float):
+        """LSTM from [..., d_model] to [..., predicted_step, num_of_vars] 
+
+        Args:
+            num_var (int): number of variables encoded in the input tensor
+            d_model (int): encoding dimension of the tensor
+            pred_step (int): step to be predicted by LSTM
+            num_layers (int): number of layers of LSTM
+            dropout (float): 
+        """
+
+        super().__init__()
+        self.num_var = num_var
+        self.d_model = d_model
+        self.num_layers = num_layers
+        self.pred_step = pred_step
+
+        self.lstm = nn.LSTM(d_model, d_model, num_layers=num_layers, batch_first=True, dropout=dropout)
+        self.linear = nn.Linear(d_model, pred_step*num_var)
+
+    def forward(self, x):
+        """LSTM process over the x tensor and reshaping according to pred_step and num_var to be predicted 
+
+        Args:
+            x (torch.Tensor): input tensor
+
+        Returns:
+            torch.Tensor: tensor resized to [B, pred_step, num_var]
+        """
+
+        h0 = torch.zeros(self.num_layers, x.size(0), self.d_model).to(x.device)
+        c0 = torch.zeros(self.num_layers, x.size(0), self.d_model).to(x.device)
+        out, _ = self.lstm(x, (h0, c0))
+        out = self.linear(out[:, -1, :])  # Take the last output of the sequence
+        out = out.view(-1, self.pred_step, self.num_var)
+        return out
+    
+class GLU(nn.Module):
+    def __init__(self, d_model: int):
+        """Gated Linear Unit
+        
+        Auxiliary subnet for sigmoid element-wise multiplication
+
+        Args:
+            d_model (int): dimension of operations
+        """
+
+        super().__init__()
+        self.linear1 = nn.Linear(d_model, d_model, bias = False)
+        self.linear2 = nn.Linear(d_model, d_model, bias = False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x1 = self.sigmoid(self.linear1(x))
+        x2 = self.linear2(x)
+        out = x1*x2 #element-wise multiplication
+        return out
+    
+class GRN(nn.Module):
+    def __init__(self, d_model: int, dropout_rate: float):
+        """Gated Residual Network
+
+        Auxiliary subnet for gating residual connections
+
+        Args:
+            d_model (int): 
+            dropout_rate (float):
+        """
+
+        super().__init__()
+        self.linear1 = nn.Linear(d_model, d_model) 
+        self.elu = nn.ELU()
+        self.linear2 = nn.Linear(d_model, d_model)
+        self.res_conn = ResidualConnection(d_model, dropout_rate)
+
+    def forward(self, x: torch.Tensor, using_norm:bool = True) -> torch.Tensor:
+        eta1 = self.elu(self.linear1(x))
+        eta2 = self.linear2(eta1)
+        out = self.res_conn(eta2, x, using_norm)
+        return out
+    
+class ResidualConnection(nn.Module):
+    def __init__(self, d_model, dropout_rate):
+        """Residual Connection of res_conn with GLU(x)
+        
+        Auxiliary subnet for residual connections
+
+        Args:
+            d_model (int): 
+            dropout_rate (float): 
+        """
+
+        super().__init__()
+        self.dropout = nn.Dropout(dropout_rate)
+        self.glu = GLU(d_model)
+        self.norm = nn.LayerNorm(d_model)
+    
+    def forward(self, x: torch.Tensor, res_conn: torch.Tensor, using_norm:bool = True) -> torch.Tensor:
+        """Res Cionnection using normalizing computatiion on 'x' and strict 'res_conn' 
+
+        Args:
+            x (torch.Tensor): GLU(dropout(x))
+            res_conn (torch.Tensor): tensor summed to x before normalization
+            using_norm (bool, optional): _description_. Defaults to True.
+
+        Returns:
+            torch.Tensor: 
+        """
+        x = self.glu(self.dropout(x))
+        out = res_conn + x
+        if using_norm:
+            out = self.norm(out)
+        return out
+
+class InterpretableMultiHead(nn.Module):
+    def __init__(self, d_model, d_head, n_head):
+        """Interpretable MultiHead Attention
+
+        Similar to canonical MultiHead Attention with Query-Keys-Value structure
+        Particularities are:
+        - Only one common "Value"-Linear layer for all heads
+        - output of all heads are summed together and then rescaled over the number of heads
+        The final output tensor is re-embedded in the initial dimension
+
+        Args:
+            d_model (int): starting and ending dimension of the net
+            d_head (int): hidden dimension of all heads
+            n_head (int): number of heads
+        """
+
+        super().__init__()
+        self.d_head = d_head
+        self.n_head = n_head
+        self.Q_layers = nn.ModuleList([nn.Linear(d_model,d_head) for _ in range(n_head)])
+        self.K_layers = nn.ModuleList([nn.Linear(d_model,d_head) for _ in range(n_head)])
+        self.Softmax_layers = nn.ModuleList([nn.Softmax(dim=-1) for _ in range(n_head)])
+        self.V_layer = nn.Linear(d_model, d_head)
+        self.out_layer = nn.Linear(d_head, d_model)
+
+    def forward(self, query:torch.Tensor, key:torch.Tensor, value:torch.Tensor) -> torch.Tensor:
+        out = torch.Tensor()
+        for (q_layer, k_layer, softmax) in zip(self.Q_layers, self.K_layers, self.Softmax_layers):
+            Q = q_layer(query)
+            K = k_layer(key)
+            wei = Q @ K.transpose(-2,-1) * (self.d_head**-0.5)
+            wei = softmax(wei)
+            V = self.V_layer(value)
+            out_h = wei @ V
+            if out.shape[0]>0:
+                out = out + out_h # sum the result of the head attention
+            else:
+                out = out_h # out is not modifies/initialized yet
+        out = out / self.n_head
+        out = self.out_layer(out) # comeback to d_model dimension
+        return out        

dsipts/models/timexer/Layers.py

@@ -0,0 +1,127 @@
+import torch.nn as nn
+import torch
+import math
+import torch.nn.functional as F
+
+
+class PositionalEmbedding(nn.Module):
+    def __init__(self, d_model, max_len=5000):
+        super(PositionalEmbedding, self).__init__()
+        # Compute the positional encodings once in log space.
+        pe = torch.zeros(max_len, d_model).float()
+        pe.require_grad = False
+
+        position = torch.arange(0, max_len).float().unsqueeze(1)
+        div_term = (torch.arange(0, d_model, 2).float()
+                    * -(math.log(10000.0) / d_model)).exp()
+
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        return self.pe[:, :x.size(1)]
+
+class EnEmbedding(nn.Module):
+    def __init__(self, n_vars, d_model, patch_len, dropout):
+        super(EnEmbedding, self).__init__()
+        # Patching
+        self.patch_len = patch_len
+
+        self.value_embedding = nn.Linear(patch_len, d_model, bias=False)
+        self.glb_token = nn.Parameter(torch.randn(1, n_vars, 1, d_model))
+        self.position_embedding = PositionalEmbedding(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        # do patching
+        n_vars = x.shape[1]
+        glb = self.glb_token.repeat((x.shape[0], 1, 1, 1))
+
+        x = x.unfold(dimension=-1, size=self.patch_len, step=self.patch_len)
+        x = torch.reshape(x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3]))
+        # Input encoding
+        x = self.value_embedding(x) + self.position_embedding(x)
+        x = torch.reshape(x, (-1, n_vars, x.shape[-2], x.shape[-1]))
+        x = torch.cat([x, glb], dim=2)
+        x = torch.reshape(x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3]))
+        return self.dropout(x), n_vars
+    
+class FlattenHead(nn.Module):
+    def __init__(self, n_vars, nf, target_window, head_dropout=0):
+        super().__init__()
+        self.n_vars = n_vars
+        self.flatten = nn.Flatten(start_dim=-2)
+        self.linear = nn.Linear(nf, target_window)
+        self.dropout = nn.Dropout(head_dropout)
+
+    def forward(self, x):  # x: [bs x nvars x d_model x patch_num]
+        x = self.flatten(x)
+        x = self.linear(x)
+        x = self.dropout(x)
+        return x
+    
+    
+
+class EncoderLayer(nn.Module):
+    def __init__(self, self_attention, cross_attention, d_model, d_ff=None,
+                 dropout=0.1, activation="relu"):
+        super(EncoderLayer, self).__init__()
+        d_ff = d_ff or 4 * d_model
+        self.self_attention = self_attention
+        self.cross_attention = cross_attention
+        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
+        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.activation = F.relu if activation == "relu" else F.gelu
+
+    def forward(self, x, cross, x_mask=None, cross_mask=None, tau=None, delta=None):
+        B, L, D = cross.shape
+        x = x + self.dropout(self.self_attention(
+            x, x, x,
+            attn_mask=x_mask,
+            tau=tau, delta=None
+        )[0])
+        x = self.norm1(x)
+
+        x_glb_ori = x[:, -1, :].unsqueeze(1)
+        x_glb = torch.reshape(x_glb_ori, (B, -1, D))
+        x_glb_attn = self.dropout(self.cross_attention(
+            x_glb, cross, cross,
+            attn_mask=cross_mask,
+            tau=tau, delta=delta
+        )[0])
+        x_glb_attn = torch.reshape(x_glb_attn,
+                                   (x_glb_attn.shape[0] * x_glb_attn.shape[1], x_glb_attn.shape[2])).unsqueeze(1)
+        x_glb = x_glb_ori + x_glb_attn
+        x_glb = self.norm2(x_glb)
+
+        y = x = torch.cat([x[:, :-1, :], x_glb], dim=1)
+
+        y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
+        y = self.dropout(self.conv2(y).transpose(-1, 1))
+
+        return self.norm3(x + y)
+    
+class Encoder(nn.Module):
+    def __init__(self, layers, norm_layer=None, projection=None):
+        super(Encoder, self).__init__()
+        self.layers = nn.ModuleList(layers)
+        self.norm = norm_layer
+        self.projection = projection
+
+    def forward(self, x, cross, x_mask=None, cross_mask=None, tau=None, delta=None):
+        for layer in self.layers:
+            x = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask, tau=tau, delta=delta)
+
+        if self.norm is not None:
+            x = self.norm(x)
+
+        if self.projection is not None:
+            x = self.projection(x)
+        return x