dsipts 1.1.5__py3-none-any.whl

dsipts/models/crossformer/cross_encoder.py

@@ -0,0 +1,99 @@
+import torch
+import torch.nn as nn
+
+from .attn import TwoStageAttentionLayer
+from math import ceil
+
+class SegMerging(nn.Module):
+    '''
+    Segment Merging Layer.
+    The adjacent `win_size' segments in each dimension will be merged into one segment to
+    get representation of a coarser scale
+    we set win_size = 2 in our paper
+    '''
+    def __init__(self, d_model, win_size, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.d_model = d_model
+        self.win_size = win_size
+        self.linear_trans = nn.Linear(win_size * d_model, d_model)
+        self.norm = norm_layer(win_size * d_model)
+
+    def forward(self, x):
+        """
+        x: B, ts_d, L, d_model
+        """
+        batch_size, ts_d, seg_num, d_model = x.shape
+        pad_num = seg_num % self.win_size
+        #import pdb
+        #pdb.set_trace()
+        if pad_num != 0: 
+            pad_num = self.win_size - pad_num
+            x = torch.cat((x, x[:, :, -pad_num:, :]), dim = -2)
+
+        seg_to_merge = []
+        for i in range(self.win_size):
+            seg_to_merge.append(x[:, :, i::self.win_size, :])
+   
+        x = torch.cat(seg_to_merge, -1)  # [B, ts_d, seg_num/win_size, win_size*d_model]
+
+        x = self.norm(x)
+        x = self.linear_trans(x)
+
+        return x
+
+class scale_block(nn.Module):
+    '''
+    We can use one segment merging layer followed by multiple TSA layers in each scale
+    the parameter `depth' determines the number of TSA layers used in each scale
+    We set depth = 1 in the paper
+    '''
+    def __init__(self, win_size, d_model, n_heads, d_ff, depth, dropout, \
+                    seg_num = 10, factor=10):
+        super(scale_block, self).__init__()
+
+        if (win_size > 1):
+            self.merge_layer = SegMerging(d_model, win_size, nn.LayerNorm)
+        else:
+            self.merge_layer = None
+        
+        self.encode_layers = nn.ModuleList()
+
+        for i in range(depth):
+            self.encode_layers.append(TwoStageAttentionLayer(seg_num, factor, d_model, n_heads, \
+                                                        d_ff, dropout))
+    
+    def forward(self, x):
+        _, ts_dim, _, _ = x.shape
+
+        if self.merge_layer is not None:
+            x = self.merge_layer(x)
+        
+        for layer in self.encode_layers:
+            x = layer(x)        
+        
+        return x
+
+class Encoder(nn.Module):
+    '''
+    The Encoder of Crossformer.
+    '''
+    def __init__(self, e_blocks, win_size, d_model, n_heads, d_ff, block_depth, dropout,
+                in_seg_num = 10, factor=10):
+        super(Encoder, self).__init__()
+        self.encode_blocks = nn.ModuleList()
+
+        self.encode_blocks.append(scale_block(1, d_model, n_heads, d_ff, block_depth, dropout,\
+                                            in_seg_num, factor))
+        for i in range(1, e_blocks):
+            self.encode_blocks.append(scale_block(win_size, d_model, n_heads, d_ff, block_depth, dropout,\
+                                            ceil(in_seg_num/win_size**i), factor))
+
+    def forward(self, x):
+        encode_x = []
+        encode_x.append(x)
+        
+        for block in self.encode_blocks:
+            x = block(x)
+            encode_x.append(x)
+
+        return encode_x

dsipts/models/d3vae/diffusion_process.py

@@ -0,0 +1,169 @@
+# -*-Encoding: utf-8 -*-
+"""
+Authors:
+    Li,Yan (liyan22021121@gmail.com)
+"""
+import numpy as np
+import torch
+from functools import partial
+from inspect import isfunction
+import torch.nn as nn
+
+
+
+
+
+def get_beta_schedule(beta_schedule, beta_start, beta_end, num_diffusion_timesteps):
+    if beta_schedule == 'quad':
+      betas = np.linspace(beta_start ** 0.5, beta_end ** 0.5, num_diffusion_timesteps, dtype=np.float64) ** 2
+    elif beta_schedule == 'linear':
+      betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64)
+    elif beta_schedule == 'const':
+      betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
+    elif beta_schedule == 'jsd':  # 1/T, 1/(T-1), 1/(T-2), ..., 1
+      betas = 1. / np.linspace(num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64)
+    else:
+      raise NotImplementedError(beta_schedule)
+    assert betas.shape == (num_diffusion_timesteps,)
+    return betas
+
+
+def default(val, d):
+    if val is not None:
+        return val
+    return d() if isfunction(d) else d
+
+
+def extract(a, t, x_shape):
+    #print(a.shape, t.shape)
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    #print(out.shape)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+
+def noise_like(shape, device, repeat=False):
+    def repeat_noise():
+        return torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
+    def noise():
+        return torch.randn(shape, device=device)
+    return repeat_noise() if repeat else noise()
+
+
+class GaussianDiffusion(nn.Module):
+    def __init__(
+        self,
+        bvae,
+        input_size,
+        beta_start=0,
+        beta_end=0.1,
+        diff_steps=100,
+        loss_type="l2",
+        betas=None,
+        scale = 0.1,
+        beta_schedule="linear",
+    ):
+        """
+        Params:
+           bave: The bidirectional vae model.
+           beta_start: The start value of the beta schedule.
+           beta_end: The end value of the beta schedule.
+           beta_schedule: the kind of the beta schedule, here are fixed to linear, you can adjust it as needed.
+           diff_steps: The maximum diffusion steps.
+           scale: scale parameters for the target time series.
+        """
+        super().__init__()
+        self.generative = bvae
+        self.scale = scale
+        self.beta_start = beta_start
+        self.beta_end = beta_end
+        betas = get_beta_schedule(beta_schedule, beta_start, beta_end, diff_steps)
+        alphas = 1.0 - betas
+        alphas_cumprod = np.cumprod(alphas, axis=0)
+    
+        alphas_target = 1.0 - betas*scale
+        alphas_target_cumprod = np.cumprod(alphas_target, axis=0)
+        self.alphas_target = alphas_target
+        self.alphas_target_cumprod = alphas_target_cumprod
+        
+        (timesteps,) = betas.shape
+        self.num_timesteps = int(timesteps)
+        self.loss_type = loss_type
+       
+        to_torch = partial(torch.tensor, dtype=torch.float32)
+        self.register_buffer("betas", to_torch(betas))
+        self.register_buffer("alphas_cumprod", to_torch(alphas_cumprod))
+    
+        self.register_buffer("sqrt_alphas_cumprod", to_torch(np.sqrt(alphas_cumprod)))
+        self.register_buffer("sqrt_alphas_target_cumprod", to_torch(np.sqrt(alphas_target_cumprod)))
+        self.register_buffer(
+            "sqrt_one_minus_alphas_cumprod", to_torch(np.sqrt(1.0 - alphas_cumprod))
+        )
+        self.register_buffer(
+            "sqrt_one_minus_alphas_target_cumprod", to_torch(np.sqrt(1.0 - alphas_target_cumprod))
+        )
+
+    def q_sample(self, x_start, t, noise=None):
+        """
+        Diffuse the initial input.
+            :param x_start: [B, T, *]
+            :return: [B, T, *]
+        """
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        return (
+            extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+            + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
+        )
+    
+    def q_sample_target(self, y_target, t, noise=None):
+        """
+        Diffuse the target.
+            :param y_target: [B1, T1, *]
+            :return: (tensor) [B1, T1, *]
+        """
+        noise = default(noise, lambda: torch.randn_like(y_target))
+       
+        return (
+            extract(self.sqrt_alphas_target_cumprod, t, y_target.shape) * y_target
+            + extract(self.sqrt_one_minus_alphas_target_cumprod, t, y_target.shape) * noise
+        )
+           
+    def p_losses(self, x_start, y_target, t,  noise=None, noise1=None):
+        """
+        Put the diffused input into the BVAE to generate the output.
+        Params
+            :param x_start: [B, T, *]
+            :param y_target: [B1, T1, *]
+            :param t: [B,]
+        -----------------------
+        Return
+            :return output: the distribution of generative results.
+            :return y_noisy: diffused target.
+            :return total_c: the total correlations of latent variables in BVAE.
+            :return all_z: all latent variables of BVAE.
+        """
+        B, T, _ = x_start.shape
+        B1, T1, _ = y_target.shape
+        x_start = x_start.reshape(B, 1, T, -1)
+        y_target = y_target.reshape(B1, 1, T1, -1)
+        
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        noise1 = default(noise1, lambda: torch.randn_like(y_target))
+
+        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise.to(x_start.device))
+       
+        y_noisy = self.q_sample_target(y_target=y_target, t=t, noise=noise1.to(y_target.device))
+        x_noisy = x_noisy.reshape(B,1, T,-1)
+        
+        y_noisy = y_noisy.reshape(B1,1, T1,-1)
+
+        logits, total_c, all_z = self.generative(x_noisy) ##forward dell'encoder
+        
+        output = self.generative.decoder_output(logits)
+        return output, y_noisy, total_c, all_z
+
+    def log_prob(self, x_input, y_target, time):
+        output, y_noisy, total_c, all_z = self.p_losses(
+            x_input, y_target, time,
+        )
+        return output, y_noisy, total_c, all_z

dsipts/models/d3vae/embedding.py

@@ -0,0 +1,108 @@
+# -*-Encoding: utf-8 -*-
+"""
+Authors:
+    Li,Yan (liyan22021121@gmail.com)
+"""
+import torch
+import torch.nn as nn
+import math
+
+
+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 TokenEmbedding(nn.Module):
+    def __init__(self, c_in, d_model):
+        super(TokenEmbedding, self).__init__()
+        padding = 1 if torch.__version__ >= '1.5.0' else 2
+        self.tokenConv = nn.Conv1d(in_channels=c_in, out_channels=d_model,
+                                   kernel_size=3, padding=padding, padding_mode='circular', bias=False)
+        for m in self.modules():
+            if isinstance(m, nn.Conv1d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='leaky_relu')
+
+    def forward(self, x):
+        x = self.tokenConv(x.permute(0, 2, 1)).transpose(1, 2)
+        return x
+
+
+class TemporalEmbedding(nn.Module):
+    def __init__(self, d_model, freq='h'):
+        super(TemporalEmbedding, self).__init__()
+
+        minute_size = 4
+        hour_size = 24
+        weekday_size = 7
+        day_size = 32
+        month_size = 13
+
+        # Embed = FixedEmbedding if embed_type=='fixed' else nn.Embedding
+        Embed = nn.Embedding
+        if freq == 't':
+            self.minute_embed = Embed(minute_size, d_model)
+            self.fc = nn.Linear(5*d_model, d_model)
+        else:
+            self.fc = nn.Linear(4*d_model, d_model)
+        self.hour_embed = Embed(hour_size, d_model)
+        self.weekday_embed = Embed(weekday_size, d_model)
+        self.day_embed = Embed(day_size, d_model)
+        self.month_embed = Embed(month_size, d_model)
+        
+    def forward(self, x):
+        x = x.long()
+        minute_x = self.minute_embed(x[:,:,4]) if hasattr(self, 'minute_embed') else 0.
+        hour_x = self.hour_embed(x[:,:,3])
+        weekday_x = self.weekday_embed(x[:,:,2])
+        day_x = self.day_embed(x[:,:,1])
+        month_x = self.month_embed(x[:,:,0])
+        if hasattr(self, 'minute_embed'):
+            out = torch.cat((minute_x, hour_x, weekday_x, day_x, month_x), dim=2)
+        else:
+            out = torch.cat((hour_x, weekday_x, day_x, month_x), dim=2)
+        # print(out.shape)
+        out = self.fc(out)
+        # print(out.shape)
+        return out
+
+
+class DataEmbedding(nn.Module):
+    def __init__(self, c_in, d_model, embs, dropout=0.1):
+        super(DataEmbedding, self).__init__()
+
+        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
+        self.position_embedding = PositionalEmbedding(d_model=d_model)
+        #self.temporal_embedding = TemporalEmbedding(d_model=d_model, freq=freq)
+        self.emb_list = nn.ModuleList()
+        if embs is not None:
+            for k in embs:
+                self.emb_list.append(nn.Embedding(k+1,d_model))
+        
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x, x_mark):
+        tot = None
+        for i in range(len(self.emb_list)):
+            if tot is None:
+                tot = self.emb_list[i](x_mark[:,:,i])
+            else:
+                tot += self.emb_list[i](x_mark[:,:,i])
+                
+        x = self.value_embedding(x) + tot + self.position_embedding(x)
+        return self.dropout(x)