Upload 4 files

model_STSSDL/STSSDL.py

@@ -0,0 +1,326 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+import math
+import numpy as np
+
+class AGCN(nn.Module):
+    def __init__(self, dim_in, dim_out, cheb_k, num_support):
+        super(AGCN, self).__init__()
+        self.cheb_k = cheb_k
+        self.weights = nn.Parameter(torch.FloatTensor(num_support*cheb_k*dim_in, dim_out)) # num_support*cheb_k*dim_in is the length of support
+        # self.weights = nn.Parameter(torch.FloatTensor(dim_in, dim_out))
+        self.bias = nn.Parameter(torch.FloatTensor(dim_out))
+        nn.init.xavier_normal_(self.weights)
+        nn.init.constant_(self.bias, val=0)
+        
+    def forward(self, x, supports):
+        x_g = []        
+        for support in supports:
+            if len(support.shape) == 2:
+                support_ks = [torch.eye(support.shape[0]).to(support.device), support]
+                for k in range(2, self.cheb_k):
+                    support_ks.append(torch.matmul(2 * support, support_ks[-1]) - support_ks[-2]) 
+                for graph in support_ks:
+                    x_g.append(torch.einsum("nm,bmc->bnc", graph, x))
+            else:
+                support_ks = [torch.eye(support.shape[1]).repeat(support.shape[0], 1, 1).to(support.device), support]
+                for k in range(2, self.cheb_k):
+                    support_ks.append(torch.matmul(2 * support, support_ks[-1]) - support_ks[-2]) 
+                for graph in support_ks:
+                    x_g.append(torch.einsum("bnm,bmc->bnc", graph, x))
+        x_g = torch.cat(x_g, dim=-1)
+        x_gconv = torch.einsum('bni,io->bno', x_g, self.weights) + self.bias  # b, N, dim_out
+        return x_gconv
+    
+class AGCRNCell(nn.Module):
+    def __init__(self, node_num, dim_in, dim_out, cheb_k, num_support):
+        super(AGCRNCell, self).__init__()
+        self.node_num = node_num
+        self.hidden_dim = dim_out
+        self.gate = AGCN(dim_in+self.hidden_dim, 2*dim_out, cheb_k, num_support)
+        self.update = AGCN(dim_in+self.hidden_dim, dim_out, cheb_k, num_support)
+
+    def forward(self, x, state, supports):
+        #x: B, num_nodes, input_dim
+        #state: B, num_nodes, hidden_dim
+        state = state.to(x.device)
+        input_and_state = torch.cat((x, state), dim=-1)
+        z_r = torch.sigmoid(self.gate(input_and_state, supports))
+        z, r = torch.split(z_r, self.hidden_dim, dim=-1)
+        candidate = torch.cat((x, z*state), dim=-1)
+        hc = torch.tanh(self.update(candidate, supports))
+        h = r*state + (1-r)*hc
+        return h
+
+    def init_hidden_state(self, batch_size):
+        return torch.zeros(batch_size, self.node_num, self.hidden_dim)
+    
+class ADCRNN_Encoder(nn.Module):
+    def __init__(self, node_num, dim_in, dim_out, cheb_k, rnn_layers, num_support):
+        super(ADCRNN_Encoder, self).__init__()
+        assert rnn_layers >= 1, 'At least one DCRNN layer in the Encoder.'
+        self.node_num = node_num
+        self.input_dim = dim_in
+        self.rnn_layers = rnn_layers
+        self.dcrnn_cells = nn.ModuleList()
+        self.dcrnn_cells.append(AGCRNCell(node_num, dim_in, dim_out, cheb_k, num_support))
+        for _ in range(1, rnn_layers):
+            self.dcrnn_cells.append(AGCRNCell(node_num, dim_out, dim_out, cheb_k, num_support))
+
+    def forward(self, x, init_state, supports):
+        #shape of x: (B, T, N, D), shape of init_state: (rnn_layers, B, N, hidden_dim)
+        assert x.shape[2] == self.node_num and x.shape[3] == self.input_dim
+        seq_length = x.shape[1]
+        current_inputs = x
+        output_hidden = []
+        for i in range(self.rnn_layers):
+            state = init_state[i]
+            inner_states = []
+            for t in range(seq_length):
+                state = self.dcrnn_cells[i](current_inputs[:, t, :, :], state, supports)
+                inner_states.append(state)
+            output_hidden.append(state)
+            current_inputs = torch.stack(inner_states, dim=1)
+        #current_inputs: the outputs of last layer: (B, T, N, hidden_dim)
+        #last_state: (B, N, hidden_dim)
+        #output_hidden: the last state for each layer: (rnn_layers, B, N, hidden_dim)
+        #return current_inputs, torch.stack(output_hidden, dim=0)
+        return current_inputs, output_hidden
+    
+    def init_hidden(self, batch_size):
+        init_states = []
+        for i in range(self.rnn_layers):
+            init_states.append(self.dcrnn_cells[i].init_hidden_state(batch_size))
+        return init_states
+
+class ADCRNN_Decoder(nn.Module):
+    def __init__(self, node_num, dim_in, dim_out, cheb_k, rnn_layers, num_support):
+        super(ADCRNN_Decoder, self).__init__()
+        assert rnn_layers >= 1, 'At least one DCRNN layer in the Decoder.'
+        self.node_num = node_num
+        self.input_dim = dim_in
+        self.rnn_layers = rnn_layers
+        self.dcrnn_cells = nn.ModuleList()
+        self.dcrnn_cells.append(AGCRNCell(node_num, dim_in, dim_out, cheb_k, num_support))
+        for _ in range(1, rnn_layers):
+            self.dcrnn_cells.append(AGCRNCell(node_num, dim_out, dim_out, cheb_k, num_support))
+
+    def forward(self, xt, init_state, supports):
+        # xt: (B, N, D)
+        # init_state: (rnn_layers, B, N, hidden_dim)
+        assert xt.shape[1] == self.node_num and xt.shape[2] == self.input_dim
+        current_inputs = xt
+        output_hidden = []
+        for i in range(self.rnn_layers):
+            state = self.dcrnn_cells[i](current_inputs, init_state[i], supports)
+            output_hidden.append(state)
+            current_inputs = state
+        return current_inputs, output_hidden
+
+
+class STSSDL(nn.Module):
+    def __init__(self, num_nodes=207, input_dim=1, output_dim=1, horizon=12, rnn_units=128, rnn_layers=1, cheb_k=3,
+                 ycov_dim=1, prototype_num=20, prototype_dim=64, tod_embed_dim=10, adj_mx=None, cl_decay_steps=2000,
+                 TDAY=288, use_curriculum_learning=True, use_STE=False, device="cpu",adaptive_embedding_dim=48,node_embedding_dim=20,input_embedding_dim=128):
+        super(STSSDL, self).__init__()
+        self.num_nodes = num_nodes
+        self.input_dim = input_dim
+        self.rnn_units = rnn_units
+        self.output_dim = output_dim
+        self.horizon = horizon
+        self.rnn_layers = rnn_layers
+        self.cheb_k = cheb_k
+        self.ycov_dim = ycov_dim
+        self.tod_embed_dim = tod_embed_dim
+        self.cl_decay_steps = cl_decay_steps
+        self.use_curriculum_learning = use_curriculum_learning
+        self.device = device
+        self.use_STE = use_STE
+        self.TDAY = TDAY
+        self.adaptive_embedding_dim=adaptive_embedding_dim
+        self.node_embedding_dim = node_embedding_dim
+        self.input_embedding_dim=input_embedding_dim
+        self.total_embedding_dim=  self.tod_embed_dim+self.adaptive_embedding_dim+self.node_embedding_dim
+        # prototypes
+        self.prototype_num = prototype_num
+        self.prototype_dim = prototype_dim
+        self.prototypes = self.construct_prototypes()
+        
+        # projection & spatio-temporal embedding
+        if self.use_STE:
+            if self.adaptive_embedding_dim > 0:
+                self.adaptive_embedding = nn.init.xavier_uniform_(
+                    nn.Parameter(torch.empty(12, num_nodes, self.adaptive_embedding_dim))
+                )
+            self.input_proj = nn.Linear(self.input_dim, input_embedding_dim)
+            self.node_embedding = nn.Parameter(torch.empty(self.num_nodes, self.node_embedding_dim))
+            self.time_embedding = nn.Parameter(torch.empty(self.TDAY, self.tod_embed_dim))
+            nn.init.xavier_uniform_(self.node_embedding)
+            nn.init.xavier_uniform_(self.time_embedding)
+
+        # encoder
+        self.adj_mx = adj_mx
+        if self.use_STE:
+            self.encoder = ADCRNN_Encoder(self.num_nodes,  input_embedding_dim + self.total_embedding_dim, self.rnn_units, self.cheb_k, self.rnn_layers, len(self.adj_mx))
+        else:
+            self.encoder = ADCRNN_Encoder(self.num_nodes, self.input_dim, self.rnn_units, self.cheb_k, self.rnn_layers, len(self.adj_mx))
+        
+        # decoder
+        self.decoder_dim = self.rnn_units + self.prototype_dim
+        if self.use_STE:
+            self.decoder = ADCRNN_Decoder(self.num_nodes, input_embedding_dim + self.total_embedding_dim-self.adaptive_embedding_dim, self.decoder_dim, self.cheb_k, self.rnn_layers, 1)
+        else:
+            self.decoder = ADCRNN_Decoder(self.num_nodes, self.output_dim + self.ycov_dim, self.decoder_dim, self.cheb_k, self.rnn_layers, 1)
+
+        # output
+        self.proj = nn.Sequential(nn.Linear(self.decoder_dim, self.output_dim, bias=True))
+        
+        # graph
+        self.hypernet = nn.Sequential(nn.Linear(self.decoder_dim*2, self.tod_embed_dim, bias=True))
+        
+        self.act_dict = {'relu': nn.ReLU(), 'lrelu': nn.LeakyReLU(), 'sigmoid': nn.Sigmoid()}
+        self.act_fn = 'sigmoid'  # 'relu' 'lrelu' 'sigmoid'
+        
+    def compute_sampling_threshold(self, batches_seen):
+        return self.cl_decay_steps / (self.cl_decay_steps + np.exp(batches_seen / self.cl_decay_steps))
+
+    def construct_prototypes(self):
+        prototypes_dict = nn.ParameterDict()
+        prototype = torch.randn(self.prototype_num, self.prototype_dim)
+        prototypes_dict['prototypes'] = nn.Parameter(prototype, requires_grad=True)     # (M, d)
+        prototypes_dict['Wq'] = nn.Parameter(torch.randn(self.rnn_units, self.prototype_dim), requires_grad=True)    # project to query
+        for param in prototypes_dict.values():
+            nn.init.xavier_normal_(param)
+
+        return prototypes_dict
+    
+    def query_prototypes(self, h_t:torch.Tensor):
+        query = torch.matmul(h_t, self.prototypes['Wq'])     # (B, N, d)
+        att_score = torch.softmax(torch.matmul(query, self.prototypes['prototypes'].t()), dim=-1)         # alpha: (B, N, M)
+        value = torch.matmul(att_score, self.prototypes['prototypes'])     # (B, N, d)
+        _, ind = torch.topk(att_score, k=2, dim=-1)
+        pos = self.prototypes['prototypes'][ind[:, :, 0]] # B, N, d
+        neg = self.prototypes['prototypes'][ind[:, :, 1]] # B, N, d
+        mask = torch.stack([ind[:, :, 0], ind[:, :, 1]], dim=-1) # B, N, 2
+            
+        return value, query, pos, neg, mask
+    
+    def calculate_distance(self, pos, pos_his, mask=None):
+        score = torch.sum(torch.abs(pos - pos_his), dim=-1)
+        return score, mask
+            
+    def forward(self, x, x_cov, x_his, y_cov, labels=None, batches_seen=None):
+        if self.use_STE:
+            if self.input_embedding_dim>0:
+                x = self.input_proj(x)  # [B,T,N,1]->[B,T,N,D]
+            features = [x]
+
+            tod = x_cov.squeeze()  # [B, T, N]
+            if self.tod_embed_dim>0:
+                time_emb = self.time_embedding[(x_cov.squeeze() * self.TDAY).type(torch.LongTensor)]  # [B, T, N, d]
+                features.append(time_emb)
+            if self.adaptive_embedding_dim > 0:
+                adp_emb = self.adaptive_embedding.expand(
+                    size=(x.shape[0], *self.adaptive_embedding.shape)
+                )
+                features.append(adp_emb)
+            if self.node_embedding_dim>0:
+                node_emb = self.node_embedding.unsqueeze(0).unsqueeze(1).expand(x.shape[0], self.horizon, -1, -1)  # [B,T,N,d]
+                features.append(node_emb)
+            x = torch.cat(features, dim=-1) # [B, T, N, D+d+80]
+        supports_en = self.adj_mx
+        init_state = self.encoder.init_hidden(x.shape[0])
+        h_en, state_en = self.encoder(x, init_state, supports_en) # B, T, N, hidden
+        h_t = h_en[:, -1, :, :] # B, N, hidden (last state)    
+        v_t, q_t, p_t, n_t, mask = self.query_prototypes(h_t)
+        if self.use_STE:
+            if self.input_embedding_dim>0:
+                x_his = self.input_proj(x_his)  # [B,T,N,1]->[B,T,N,D]
+            features = [x_his]
+            tod = x_cov.squeeze()  # [B, T, N]
+            if self.tod_embed_dim>0:
+                time_emb = self.time_embedding[(x_cov.squeeze() * self.TDAY).type(torch.LongTensor)]  # [B, T, N, d]
+
+                features.append(time_emb)
+            if self.adaptive_embedding_dim > 0:
+                adp_emb = self.adaptive_embedding.expand(
+                        size=(x.shape[0], *self.adaptive_embedding.shape)
+                    )
+                features.append(adp_emb)
+            if self.node_embedding_dim>0:
+                node_emb = self.node_embedding.unsqueeze(0).unsqueeze(1).expand(x.shape[0], self.horizon, -1, -1)  # [B,T,N,d]
+                features.append(node_emb)
+            x_his = torch.cat(features, dim=-1) # [B, T, N, D+d+80]
+        h_his_en, state_his_en = self.encoder(x_his, init_state, supports_en) # B, T, N, hidden
+        h_a = h_his_en[:, -1, :, :] # B, N, hidden (last state)      
+        v_a, q_a, p_a, n_a, mask_his = self.query_prototypes(h_a)
+        
+        latent_dis, _ = self.calculate_distance(q_t, q_a)
+        prototype_dis, mask_dis = self.calculate_distance(p_t, p_a)
+        
+        query = torch.stack([q_t, q_a], dim=0)
+        pos = torch.stack([p_t, p_a], dim=0)
+        neg = torch.stack([n_t, n_a], dim=0)
+        mask = torch.stack([mask, mask_his], dim=0) if mask is not None else [None, None] 
+        
+        h_de = torch.cat([h_t, v_t], dim=-1)
+        h_aug = torch.cat([h_t, v_t, h_a, v_a], dim=-1) # B, N, D
+        
+        node_embeddings = self.hypernet(h_aug) # B, N, e
+        support = F.softmax(F.relu(torch.einsum('bnc,bmc->bnm', node_embeddings, node_embeddings)), dim=-1) 
+        supports_de = [support]
+        
+        ht_list = [h_de]*self.rnn_layers
+        go = torch.zeros((x.shape[0], self.num_nodes, self.output_dim), device=x.device)
+        
+        out = []
+        for t in range(self.horizon):
+            if self.use_STE:
+                if self.input_embedding_dim>0:
+                    go = self.input_proj(go)  # equal to torch.zeros(B,N,D)
+                features = [go]
+                tod = y_cov[:, t, ...].squeeze()  # [B, T, N]
+                if self.tod_embed_dim>0:
+                    time_emb = self.time_embedding[(tod * self.TDAY).type(torch.LongTensor)]
+                    features.append(time_emb)
+                if self.node_embedding_dim>0:
+                    node_emb = self.node_embedding.unsqueeze(0).expand(x.shape[0], -1, -1)  # [B,N,d]
+                    features.append(node_emb)
+                go = torch.cat(features, dim=-1) # [B, T, N, D+d]
+                h_de, ht_list = self.decoder(go, ht_list, supports_de)
+            else:
+                h_de, ht_list = self.decoder(torch.cat([go, y_cov[:, t, ...]], dim=-1), ht_list, supports_de)
+            go = self.proj(h_de)
+            out.append(go)
+            if self.training and self.use_curriculum_learning:
+                c = np.random.uniform(0, 1)
+                if c < self.compute_sampling_threshold(batches_seen):
+                    go = labels[:, t, ...]
+                    
+        output = torch.stack(out, dim=1)
+        
+        return output, query, pos, neg, mask, latent_dis, prototype_dis
+
+def print_params(model):
+    # print trainable params
+    param_count = 0
+    print('Trainable parameter list:')
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            print(name, param.shape, param.numel())
+            param_count += param.numel()
+    print(f'In total: {param_count} trainable parameters.')
+    return
+
+def main():
+    from torchinfo import summary
+    from utils import load_adj
+    
+    adj_mx = load_adj('../METRLA/adj_mx.pkl', "symadj")
+    adj_mx = [torch.FloatTensor(i) for i in adj_mx]
+    model = STSSDL(adj_mx=adj_mx)
+    summary(model, [[8, 12, 207, 1], [8, 12, 207, 1], [8, 12, 207, 1], [8, 12, 207, 1]], device="cpu")
+    
+if __name__ == '__main__':
+    main()

model_STSSDL/metrics.py

@@ -0,0 +1,38 @@
+import torch
+
+def MSE(y_true, y_pred, mask=None):
+    
+    mse = torch.square(y_pred - y_true)
+    if mask is not None:
+        mse_ig = mse * (1 - mask)  # max: 0.5 (0.9422?) min:0, but max should be close 0
+        mse = mse * mask  # max: 0.65, min:0
+        anomoly_num = torch.sum(mask)
+        normal_num = torch.sum(1 - mask)  # in general, normal = 3 * anomoly
+        mse_ig = torch.mean(mse_ig)  # mae_ig = 3 * mae alought max:0.9422, mean:0.04, i.e., most is accurately predicted
+    mse = torch.mean(mse)
+    
+    return mse
+    
+def RMSE(y_true, y_pred, mask=None):
+    
+    rmse = torch.square(torch.abs(y_pred - y_true))
+    rmse = torch.sqrt(MSE(y_true, y_pred, mask))
+    return rmse
+        
+def MAE(y_true, y_pred, mask=None):
+    
+    mae = torch.abs(y_pred - y_true)
+    if mask is not None:
+        mae_ig = mae * (1 - mask)  # max: 0.5 (0.9422?) min:0, but max should be close 0
+        mae = mae * mask  # max: 0.65, min:0
+        anomoly_num = torch.sum(mask)
+        normal_num = torch.sum(1 - mask)  # in general, normal = 3 * anomoly
+        mae_ig = torch.mean(mae_ig)  # mae_ig = 3 * mae alought max:0.9422, mean:0.04, i.e., most is accurately predicted
+    mae = torch.mean(mae)
+
+    return mae
+
+
+
+
+

model_STSSDL/train_STSSDL.py

@@ -0,0 +1,442 @@
+import sys
+import os
+import shutil
+import numpy as np
+import pandas as pd
+import time
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+import torch.nn.functional as F
+from torchinfo import summary
+import argparse
+import logging
+from utils import StandardScaler, masked_mae_loss, masked_mape_loss, masked_mse_loss, masked_rmse_loss
+from utils import load_adj
+from metrics import RMSE, MAE, MSE
+from STSSDL import STSSDL
+import random
+class ContrastiveLoss():
+    def __init__(self, contra_loss='triplet', mask=None, temp=1.0, margin=0.5):
+        self.infonce = contra_loss in ['infonce']
+        self.mask = mask
+        self.temp = temp
+        self.margin = margin
+    
+    def calculate(self, query, pos, neg, mask):
+        """
+        :param query: shape (batch_size, num_sensor, hidden_dim)
+        :param pos: shape (batch_size, num_sensor, hidden_dim)
+        :param neg: shape (batch_size, num_sensor, hidden_dim) or (batch_size, num_sensor, num_prototypes, hidden_dim)
+        :param mask: shape (batch_size, num_sensor, num_prototypes) True means positives
+        """
+        contrastive_loss = nn.TripletMarginLoss(margin=self.margin)
+        return contrastive_loss(query.detach(), pos, neg)
+
+
+def print_model(model):
+    param_count = 0
+    logger.info('Trainable parameter list:')
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            print(name, param.shape, param.numel())
+            param_count += param.numel()
+    logger.info(f'In total: {param_count} trainable parameters.')
+    return
+
+def get_model():
+    adj_mx = load_adj(adj_mx_path, args.adj_type)
+    adjs = [torch.tensor(i).to(device) for i in adj_mx]            
+    model = STSSDL(num_nodes=args.num_nodes, input_dim=args.input_dim, output_dim=args.output_dim, horizon=args.horizon, 
+                 rnn_units=args.rnn_units, rnn_layers=args.rnn_layers, cheb_k = args.cheb_k, prototype_num=args.prototype_num, 
+                 prototype_dim=args.prototype_dim, tod_embed_dim=args.tod_embed_dim, adj_mx = adjs, cl_decay_steps=args.cl_decay_steps, 
+                 use_curriculum_learning=args.use_curriculum_learning, use_STE=args.use_STE, adaptive_embedding_dim=args.adaptive_embedding_dim,node_embedding_dim=args.node_embedding_dim,input_embedding_dim=args.input_embedding_dim,device=device).to(device)
+    return model
+
+def prepare_x_y(x, y):
+    """
+    :param x: shape (batch_size, seq_len, num_sensor, input_dim)
+    :param y: shape (batch_size, horizon, num_sensor, input_dim)
+    :return1: x shape (seq_len, batch_size, num_sensor, input_dim)
+              y shape (horizon, batch_size, num_sensor, input_dim)
+    :return2: x: shape (seq_len, batch_size, num_sensor * input_dim)
+              y: shape (horizon, batch_size, num_sensor * output_dim)
+    """
+    x0 = x[..., 0:1]
+    x1 = x[..., 1:2]
+    x2 = x[..., 2:3]  
+    y0 = y[..., 0:1]
+    y1 = y[..., 1:2]
+    return x0, x1, x2, y0, y1 # x, x_cov, x_his, y, y_cov
+
+def evaluate(model, mode):
+    with torch.no_grad():
+        model = model.eval()
+        data_iter =  data[f'{mode}_loader']
+        ys_true, ys_pred = [], []
+        losses = []
+        for x, y in data_iter:
+            x = x.to(device)
+            y = y.to(device)
+            x, x_cov, x_his, y, y_cov = prepare_x_y(x, y)
+            output, _, _, _, _, _, _ = model(x, x_cov, x_his, y_cov)
+            y_pred = scaler.inverse_transform(output)
+            y_true = y
+            ys_true.append(y_true)
+            ys_pred.append(y_pred)
+            losses.append(masked_mae_loss(y_pred, y_true).item())
+        
+        ys_true, ys_pred = torch.cat(ys_true, dim=0), torch.cat(ys_pred, dim=0)
+        loss = masked_mae_loss(ys_pred, ys_true)
+
+        if mode == 'test':
+            mae = masked_mae_loss(ys_pred, ys_true).item()
+            mape = masked_mape_loss(ys_pred, ys_true).item()
+            rmse = masked_rmse_loss(ys_pred, ys_true).item()
+            mae_3 = masked_mae_loss(ys_pred[:, 2, ...], ys_true[:, 2, ...]).item()
+            mape_3 = masked_mape_loss(ys_pred[:, 2, ...], ys_true[:, 2, ...]).item()
+            rmse_3 = masked_rmse_loss(ys_pred[:, 2, ...], ys_true[:, 2, ...]).item()
+            mae_6 = masked_mae_loss(ys_pred[:, 5, ...], ys_true[:, 5, ...]).item()
+            mape_6 = masked_mape_loss(ys_pred[:, 5, ...], ys_true[:, 5, ...]).item()
+            rmse_6 = masked_rmse_loss(ys_pred[:, 5, ...], ys_true[:, 5, ...]).item()
+            mae_12 = masked_mae_loss(ys_pred[:, 11, ...], ys_true[:, 11, ...]).item()
+            mape_12 = masked_mape_loss(ys_pred[:, 11, ...], ys_true[:, 11, ...]).item()
+            rmse_12 = masked_rmse_loss(ys_pred[:, 11, ...], ys_true[:, 11, ...]).item()
+            
+            logger.info('Horizon overall: mae: {:.4f}, mape: {:.4f}, rmse: {:.4f}'.format(mae, mape * 100, rmse))
+            logger.info('Horizon 15mins: mae: {:.4f}, mape: {:.4f}, rmse: {:.4f}'.format(mae_3, mape_3 * 100, rmse_3))
+            logger.info('Horizon 30mins: mae: {:.4f}, mape: {:.4f}, rmse: {:.4f}'.format(mae_6, mape_6 * 100, rmse_6))
+            logger.info('Horizon 60mins: mae: {:.4f}, mape: {:.4f}, rmse: {:.4f}'.format(mae_12, mape_12 * 100, rmse_12))
+
+        return np.mean(losses), ys_true, ys_pred
+
+    
+def traintest_model():  
+    model = get_model()
+    print_model(model)
+    
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, eps=args.epsilon, weight_decay=args.weight_decay)
+    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.steps, gamma=args.lr_decay_ratio)
+    min_val_loss = float('inf')
+    wait = 0
+    batches_seen = 0
+    for epoch_num in range(args.epochs):
+        start_time = time.time()
+        model = model.train()
+        data_iter = data['train_loader']
+        losses, mae_losses, contra_losses, deviation_losses = [], [], [], []
+        for x, y in data_iter:
+            optimizer.zero_grad()
+            x = x.to(device)
+            y = y.to(device)
+            x, x_cov, x_his, y, y_cov = prepare_x_y(x, y)
+            output, query, pos, neg, mask, query_simi, pos_simi = model(x, x_cov, x_his, y_cov, scaler.transform(y), batches_seen)
+            y_pred = scaler.inverse_transform(output)
+            y_true = y
+
+            mae_loss = masked_mae_loss(y_pred, y_true) # masked_mae_loss(y_pred, y_true)
+            contrastive_loss = ContrastiveLoss(contra_loss=args.contra_loss, mask=mask, temp=args.temp)
+            
+            loss_c = contrastive_loss.calculate(query[0], pos[0], neg[0], mask[0])
+            loss_d = F.l1_loss(query_simi.detach(), pos_simi)
+            loss = mae_loss + args.lamb_c * loss_c + args.lamb_d * loss_d
+            
+            losses.append(loss.item())
+            mae_losses.append(mae_loss.item())
+            contra_losses.append(loss_c.item())
+            deviation_losses.append(loss_d.item())
+            losses.append(loss.item())
+            batches_seen += 1
+            loss.backward()
+            if args.max_grad_norm:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm) # gradient clipping - this does it in place
+            optimizer.step()
+            
+        end_time2 = time.time()
+        train_loss = np.mean(losses)
+        train_mae_loss = np.mean(mae_losses) 
+        train_contra_loss = np.mean(contra_losses)
+        train_deviation_loss = np.mean(deviation_losses)
+        lr_scheduler.step()
+        val_loss, _, _ = evaluate(model, 'val')
+        message = 'Epoch [{}/{}] ({}) train_loss: {:.4f}, train_mae_loss: {:.4f}, train_contra_loss: {:.4f}, train_deviation_loss: {:.4f}, val_loss: {:.4f}, lr: {:.6f}, {:.2f}s'.format(epoch_num + 1, args.epochs, batches_seen, train_loss, train_mae_loss, train_contra_loss, train_deviation_loss, val_loss, optimizer.param_groups[0]['lr'], (end_time2 - start_time))
+        logger.info(message)
+        
+        test_loss, _, _ = evaluate(model, 'test')
+        logger.info("\n")
+        
+        if val_loss < min_val_loss:
+            wait = 0
+            min_val_loss = val_loss
+            torch.save(model.state_dict(), modelpt_path)
+        elif val_loss >= min_val_loss:
+            wait += 1
+            if wait == args.patience:
+                logger.info('Early stopping at epoch: %d' % (epoch_num + 1))
+                break
+    
+    logger.info('=' * 35 + 'Best val_loss model performance' + '=' * 35)
+    logger.info('=' * 22 + 'Better results might be found from model at different epoch' + '=' * 22)
+    model = get_model()
+    model.load_state_dict(torch.load(modelpt_path))
+    start=time.time()
+    test_loss, _, _ = evaluate(model, 'test')
+    end=time.time()
+    logger.info(f"Inference Time: {(end-start):.2f}s")
+
+#########################################################################################    
+parser = argparse.ArgumentParser()
+parser.add_argument('--dataset', type=str, choices=['METRLA', 'PEMSBAY','PEMS04','PEMS07','PEMS08','PEMSD7M'], default='METRLA', help='which dataset to run')
+parser.add_argument('--num_nodes', type=int, default=207, help='num_nodes')
+parser.add_argument('--seq_len', type=int, default=12, help='input sequence length')
+parser.add_argument('--horizon', type=int, default=12, help='output sequence length')
+parser.add_argument('--input_dim', type=int, default=1, help='number of input channel')
+parser.add_argument('--output_dim', type=int, default=1, help='number of output channel')
+parser.add_argument('--tod_embed_dim', type=int, default=10, help='embedding dimension for adaptive graph')
+parser.add_argument('--cheb_k', type=int, default=3, help='max diffusion step or Cheb K')
+parser.add_argument('--rnn_layers', type=int, default=1, help='number of rnn layers')
+parser.add_argument('--rnn_units', type=int, default=128, help='number of rnn units')
+parser.add_argument('--prototype_num', type=int, default=20, help='number of meta-nodes/prototypes')
+parser.add_argument('--prototype_dim', type=int, default=64, help='dimension of meta-nodes/prototypes')
+parser.add_argument("--loss", type=str, default='mask_mae_loss', help="mask_mae_loss")
+parser.add_argument("--epochs", type=int, default=200, help="number of epochs of training")
+parser.add_argument("--patience", type=int, default=30, help="patience used for early stop")
+parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
+parser.add_argument("--lr", type=float, default=0.01, help="base learning rate")
+parser.add_argument("--steps", type=eval, default=[50, 100], help="steps") 
+parser.add_argument("--lr_decay_ratio", type=float, default=0.1, help="lr_decay_ratio")
+parser.add_argument("--weight_decay", type=float, default=0, help="weight_decay_ratio")
+parser.add_argument("--epsilon", type=float, default=1e-3, help="optimizer epsilon")
+parser.add_argument("--max_grad_norm", type=int, default=5, help="max_grad_norm")
+parser.add_argument("--use_curriculum_learning", type=eval, choices=[True, False], default='True', help="use_curriculum_learning")
+parser.add_argument("--adj_type", type=str, default='symadj', help="scalap, normlap, symadj, transition, doubletransition")
+parser.add_argument("--cl_decay_steps", type=int, default=2000, help="cl_decay_steps")
+parser.add_argument('--gpu', type=int, default=0, help='which gpu to use')
+parser.add_argument('--seed', type=int, default=100, help='random seed.')
+parser.add_argument('--temp', type=float, default=1.0, help='temperature parameter')
+parser.add_argument('--lamb_c', type=float, default=0.1, help='contra loss lambda') 
+parser.add_argument('--lamb_d', type=float, default=1.0, help='deviation loss lambda')
+parser.add_argument('--contra_loss', type=str, choices=['triplet', 'infonce'], default='triplet', help='whether to triplet or infonce contra loss')
+parser.add_argument("--use_STE", type=eval, choices=[True, False], default='True', help="use spatio-temporal embedding")
+parser.add_argument("--adaptive_embedding_dim", type=int,default=48, help="use spatio-temporal adaptive embedding")
+parser.add_argument("--node_embedding_dim", type=int,default=20, help="use spatio-temporal adaptive embedding")
+parser.add_argument("--input_embedding_dim", type=int,default=128, help="use spatio-temporal adaptive embedding")
+
+args = parser.parse_args()
+num_nodes_dict={
+    "METRLA": 207,
+    "PEMSBAY": 325,
+    "PEMS04": 307,
+    "PEMS07": 883,
+    "PEMS08": 170,
+    "PEMSD7M": 228,
+}
+if args.dataset == 'METRLA':
+    data_path = f'../{args.dataset}/metr-la.h5'
+    adj_mx_path = f'../{args.dataset}/adj_mx.pkl'
+    args.num_nodes = 207
+    args.use_STE=True
+    rand_seed=random.randint(0, 1000000)# 31340
+    args.seed=999
+    args.lamb_c=0.01
+    args.lamb_d=1
+    args.steps = [50,70]
+    args.input_embedding_dim=3
+    args.node_embedding_dim=25
+    args.tod_embed_dim=20 #TOD embedding
+    args.adaptive_embedding_dim=0
+    
+elif args.dataset == 'PEMSBAY':
+    data_path = f'../{args.dataset}/pems-bay.h5'
+    adj_mx_path = f'../{args.dataset}/adj_mx_bay.pkl'
+    args.num_nodes = 325
+    args.use_STE=True
+    args.cl_decay_steps = 8000
+    args.steps = [10, 70,150]
+    args.seed=666
+    args.lamb_c=0.01
+    args.lamb_d=1
+    args.input_embedding_dim=10
+    args.node_embedding_dim=20
+    args.tod_embed_dim=20 #TOD embedding
+    args.adaptive_embedding_dim=0
+
+elif args.dataset == 'PEMS04':
+    data_path = f'../{args.dataset}/{args.dataset}.npz'
+    adj_mx_path = f'../{args.dataset}/adj_{args.dataset}_distance.pkl'
+    args.num_nodes = num_nodes_dict[args.dataset]
+    rand_seed=random.randint(0, 1000000)# 31340
+    args.seed=610958
+    args.patience=30
+    args.batch_size=16
+    args.lr=0.001
+    args.epochs=200
+    args.steps=[50, 100]
+    args.weight_decay=0
+    args.max_grad_norm=0
+    args.rnn_units=32
+    args.prototype_num=20
+    args.prototype_dim=64
+    args.cl_decay_steps=6000
+    args.max_diffusion_step=3
+    args.input_embedding_dim=32
+    args.node_embedding_dim=24
+    args.tod_embed_dim=40 #TOD embedding
+    args.adaptive_embedding_dim=0
+    args.use_curriculum_learning=True
+    args.lamb_c=0.01
+    args.lamb_d=0.01
+
+    
+elif args.dataset == 'PEMS07':
+    data_path = f'../{args.dataset}/{args.dataset}.npz'
+    adj_mx_path = f'../{args.dataset}/adj_{args.dataset}_distance.pkl'
+    args.num_nodes = num_nodes_dict[args.dataset]
+    args.patience=20
+    args.batch_size=16
+    args.lr=0.001
+    args.steps=[50, 100]
+    args.weight_decay=0
+    args.max_grad_norm=0
+    args.rnn_units=64
+    args.prototype_num=20
+    args.prototype_dim=64
+    args.cl_decay_steps=6000
+    args.max_diffusion_step=3
+    args.lamb_c=0.01
+    args.lamb_d=1
+    args.seed=100
+    args.input_embedding_dim=64
+    args.node_embedding_dim=16
+    args.tod_embed_dim=16 
+    args.adaptive_embedding_dim=0
+elif args.dataset == 'PEMS08':
+    data_path = f'../{args.dataset}/{args.dataset}.npz'
+    adj_mx_path = f'../{args.dataset}/adj_{args.dataset}_distance.pkl'
+    args.num_nodes = num_nodes_dict[args.dataset]
+    args.use_STE=True
+    args.patience=20
+    args.batch_size=16
+    rand_seed=random.randint(0, 1000000)# 31340
+    args.seed=rand_seed
+    args.cl_decay_steps=6000
+    args.max_diffusion_step=3
+    args.steps=[70, 100]
+    args.prototype_num=20
+    args.prototype_dim=64
+    args.use_curriculum_learning=True
+    args.rnn_units = 12 
+    args.lamb_c=0.1
+    args.lamb_d=1
+    args.input_embedding_dim=16
+    args.node_embedding_dim=20
+    args.tod_embed_dim=20 #TOD embedding
+    args.adaptive_embedding_dim=0
+    
+elif args.dataset == 'PEMSD7M':
+    data_path = f'../{args.dataset}/{args.dataset}.npz'
+    adj_mx_path = f'../{args.dataset}/adj_{args.dataset}_distance.pkl'
+    args.num_nodes = num_nodes_dict[args.dataset]
+    rand_seed=random.randint(0, 1000000)# 31340
+    args.seed=119089
+    args.patience=30
+    args.batch_size=16
+    args.lr=0.001
+    args.steps=[50, 100]
+    args.weight_decay=0
+    args.max_grad_norm=0
+    args.rnn_units=32
+    args.prototype_num=16
+    args.prototype_dim=64
+    args.cl_decay_steps=4000
+    args.max_diffusion_step=3
+    args.lamb_c=0.1
+    args.lamb_d=1
+    args.input_embedding_dim=32
+    args.node_embedding_dim=20
+    args.tod_embed_dim=16 #TOD embedding
+    args.adaptive_embedding_dim=0
+    
+model_name = 'STSSDL'
+timestring = time.strftime('%Y%m%d%H%M%S', time.localtime())
+path = f'../save/{args.dataset}_{model_name}_{timestring}'
+logging_path = f'{path}/{model_name}_{timestring}_logging.txt'
+score_path = f'{path}/{model_name}_{timestring}_scores.txt'
+epochlog_path = f'{path}/{model_name}_{timestring}_epochlog.txt'
+modelpt_path = f'{path}/{model_name}_{timestring}.pt'
+if not os.path.exists(path): os.makedirs(path)
+shutil.copy2(sys.argv[0], path)
+shutil.copy2(f'{model_name}.py', path)
+shutil.copy2('utils.py', path)
+    
+logger = logging.getLogger(__name__)
+logger.setLevel(level = logging.INFO)
+class MyFormatter(logging.Formatter):
+    def format(self, record):
+        spliter = ' '
+        record.msg = str(record.msg) + spliter + spliter.join(map(str, record.args))
+        record.args = tuple() # set empty to args
+        return super().format(record)
+formatter = MyFormatter()
+handler = logging.FileHandler(logging_path, mode='a')
+handler.setLevel(logging.INFO)
+handler.setFormatter(formatter)
+console = logging.StreamHandler()
+console.setLevel(logging.INFO)
+console.setFormatter(formatter)
+logger.addHandler(handler)
+logger.addHandler(console)
+message = ''.join([f'{k}: {v}\n' for k, v in vars(args).items()])
+logger.info(message)
+
+cpu_num = 1
+os.environ ['OMP_NUM_THREADS'] = str(cpu_num)
+os.environ ['OPENBLAS_NUM_THREADS'] = str(cpu_num)
+os.environ ['MKL_NUM_THREADS'] = str(cpu_num)
+os.environ ['VECLIB_MAXIMUM_THREADS'] = str(cpu_num)
+os.environ ['NUMEXPR_NUM_THREADS'] = str(cpu_num)
+torch.set_num_threads(cpu_num)
+device = torch.device("cuda:{}".format(args.gpu)) if torch.cuda.is_available() else torch.device("cpu")
+
+np.random.seed(args.seed)
+torch.manual_seed(args.seed)
+if torch.cuda.is_available(): torch.cuda.manual_seed(args.seed)
+#####################################################################################################
+
+data = {}
+for category in ['train', 'val', 'test']:
+    cat_data = np.load(os.path.join(f'../{args.dataset}', category + 'his.npz'))
+    data['x_' + category] = np.nan_to_num(cat_data['x']) if True in np.isnan(cat_data['x']) else cat_data['x']
+    data['y_' + category] = np.nan_to_num(cat_data['y']) if True in np.isnan(cat_data['y']) else cat_data['y']
+scaler = StandardScaler(mean=data['x_train'][..., 0].mean(), std=data['x_train'][..., 0].std())
+for category in ['train', 'val', 'test']:
+    data['x_' + category][..., 0] = scaler.transform(data['x_' + category][..., 0])
+    data['x_' + category][..., 2] = scaler.transform(data['x_' + category][..., 2]) # x_his
+
+data['train_loader'] = torch.utils.data.DataLoader(
+    torch.utils.data.TensorDataset(torch.FloatTensor(data['x_train']), torch.FloatTensor(data['y_train'])),
+    batch_size=args.batch_size,
+    shuffle=True
+)
+data['val_loader'] = torch.utils.data.DataLoader(
+    torch.utils.data.TensorDataset(torch.FloatTensor(data['x_val']), torch.FloatTensor(data['y_val'])),
+    batch_size=args.batch_size, 
+    shuffle=False
+)
+data['test_loader'] = torch.utils.data.DataLoader(
+    torch.utils.data.TensorDataset(torch.FloatTensor(data['x_test']), torch.FloatTensor(data['y_test'])),
+    batch_size=args.batch_size, 
+    shuffle=False
+)
+
+def main():
+    logger.info(args.dataset, 'training and testing started', time.ctime())
+    logger.info('train xs.shape, ys.shape', data['x_train'].shape, data['y_train'].shape)
+    logger.info('val xs.shape, ys.shape', data['x_val'].shape, data['y_val'].shape)
+    logger.info('test xs.shape, ys.shape', data['x_test'].shape, data['y_test'].shape)
+    traintest_model()
+    logger.info(args.dataset, 'training and testing ended', time.ctime())
+    
+if __name__ == '__main__':
+    main()
+    

model_STSSDL/utils.py

@@ -0,0 +1,264 @@
+import pickle
+import torch
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+from scipy.sparse import linalg
+
+class DataLoader(object):
+    def __init__(self, xs, ys, batch_size, pad_with_last_sample=True, shuffle=False):
+        """
+
+        :param xs:
+        :param ys:
+        :param batch_size:
+        :param pad_with_last_sample: pad with the last sample to make number of samples divisible to batch_size.
+        """
+        self.batch_size = batch_size
+        self.current_ind = 0
+        if pad_with_last_sample:
+            num_padding = (batch_size - (len(xs) % batch_size)) % batch_size
+            x_padding = np.repeat(xs[-1:], num_padding, axis=0)
+            y_padding = np.repeat(ys[-1:], num_padding, axis=0)
+            xs = np.concatenate([xs, x_padding], axis=0)
+            ys = np.concatenate([ys, y_padding], axis=0)
+        self.size = len(xs)
+        self.num_batch = int(self.size // self.batch_size)
+        if shuffle:
+            permutation = np.random.permutation(self.size)
+            xs, ys = xs[permutation], ys[permutation]
+        self.xs = xs
+        self.ys = ys
+
+    def get_iterator(self):
+        self.current_ind = 0
+
+        def _wrapper():
+            while self.current_ind < self.num_batch:
+                start_ind = self.batch_size * self.current_ind
+                end_ind = min(self.size, self.batch_size * (self.current_ind + 1))
+                x_i = self.xs[start_ind: end_ind, ...]
+                y_i = self.ys[start_ind: end_ind, ...]
+                yield (x_i, y_i)
+                self.current_ind += 1
+
+        return _wrapper()
+    
+class StandardScaler():
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def transform(self, data):
+        return (data - self.mean) / self.std
+
+    def inverse_transform(self, data):
+        return (data * self.std) + self.mean
+    
+def getTimestamp(data):
+    num_samples, num_nodes = data.shape
+    time_ind = (data.index.values - data.index.values.astype("datetime64[D]")) / np.timedelta64(1, "D")
+    time_in_day = np.tile(time_ind, [num_nodes,1]).transpose((1, 0))
+    return time_in_day
+
+def getDayTimestamp(data):
+    # 288 timeslots each day for dataset has 5 minutes time interval.
+    df = pd.DataFrame({'timestamp':data.index.values})
+    df['weekdaytime'] = df['timestamp'].dt.weekday * 288 + (df['timestamp'].dt.hour * 60 + df['timestamp'].dt.minute)//5
+    df['weekdaytime'] = df['weekdaytime'] / df['weekdaytime'].max()
+    num_samples, num_nodes = data.shape
+    time_ind = df['weekdaytime'].values
+    time_ind_node = np.tile(time_ind, [num_nodes,1]).transpose((1, 0))
+    return time_ind_node
+
+def getDayTimestamp_(start, end, freq, num_nodes):
+    # 288 timeslots each day for dataset has 5 minutes time interval.
+    df = pd.DataFrame({'timestamp':pd.date_range(start=start, end=end, freq=freq)})
+    df['weekdaytime'] = df['timestamp'].dt.weekday * 288 + (df['timestamp'].dt.hour * 60 + df['timestamp'].dt.minute)//5
+    df['weekdaytime'] = df['weekdaytime'] / df['weekdaytime'].max()
+    time_ind = df['weekdaytime'].values
+    time_ind_node = np.tile(time_ind, [num_nodes, 1]).transpose((1, 0))
+    return time_ind_node
+
+def masked_mse(preds, labels, null_val=1e-3):
+    if np.isnan(null_val):
+        mask = ~torch.isnan(labels)
+    else:
+        mask = (labels > null_val)
+    mask = mask.float()
+    mask /= torch.mean((mask))
+    mask = torch.where(torch.isnan(mask), torch.zeros_like(mask), mask)
+    loss = (preds-labels)**2
+    loss = loss * mask
+    loss = torch.where(torch.isnan(loss), torch.zeros_like(loss), loss)
+    return torch.mean(loss)
+
+def masked_rmse(preds, labels, null_val=1e-3):
+    return torch.sqrt(masked_mse(preds=preds, labels=labels, null_val=null_val))
+
+
+def masked_mae(preds, labels, null_val=1e-3):
+    if np.isnan(null_val):
+        mask = ~torch.isnan(labels)
+    else:
+        mask = (labels > null_val)
+    mask = mask.float()
+    mask /=  torch.mean((mask))
+    mask = torch.where(torch.isnan(mask), torch.zeros_like(mask), mask)
+    loss = torch.abs(preds-labels)
+    loss = loss * mask
+    loss = torch.where(torch.isnan(loss), torch.zeros_like(loss), loss)
+    return torch.mean(loss)
+
+
+def masked_mape(preds, labels, null_val=1e-3):
+    if np.isnan(null_val):
+        mask = ~torch.isnan(labels)
+    else:
+        mask = (labels > null_val)
+    mask = mask.float()
+    mask /=  torch.mean((mask))
+    mask = torch.where(torch.isnan(mask), torch.zeros_like(mask), mask)
+    loss = torch.abs(preds-labels)/labels
+    loss = loss * mask
+    loss = torch.where(torch.isnan(loss), torch.zeros_like(loss), loss)
+    return torch.mean(loss)
+
+# DCRNN
+def masked_mae_loss(y_pred, y_true):
+    mask = (y_true != 0).float()
+    mask /= mask.mean()
+    loss = torch.abs(y_pred - y_true)
+    loss = loss * mask
+    # trick for nans: https://discuss.pytorch.org/t/how-to-set-nan-in-tensor-to-0/3918/3
+    loss[loss != loss] = 0
+    return loss.mean()
+
+def masked_mape_loss(y_pred, y_true):
+    mask = (y_true != 0).float()
+    mask /= mask.mean()
+    loss = torch.abs(torch.div(y_true - y_pred, y_true))
+    loss = loss * mask
+    # trick for nans: https://discuss.pytorch.org/t/how-to-set-nan-in-tensor-to-0/3918/3
+    loss[loss != loss] = 0
+    return loss.mean()
+
+def masked_rmse_loss(y_pred, y_true):
+    mask = (y_true != 0).float()
+    mask /= mask.mean()
+    loss = torch.pow(y_true - y_pred, 2)
+    loss = loss * mask
+    # trick for nans: https://discuss.pytorch.org/t/how-to-set-nan-in-tensor-to-0/3918/3
+    loss[loss != loss] = 0
+    return torch.sqrt(loss.mean())
+
+def masked_mse_loss(y_pred, y_true):
+    mask = (y_true != 0).float()
+    mask /= mask.mean()
+    loss = torch.pow(y_true - y_pred, 2)
+    loss = loss * mask
+    # trick for nans: https://discuss.pytorch.org/t/how-to-set-nan-in-tensor-to-0/3918/3
+    loss[loss != loss] = 0
+    return loss.mean()
+
+def load_pickle(pickle_file):
+    try:
+        with open(pickle_file, 'rb') as f:
+            pickle_data = pickle.load(f)
+    except UnicodeDecodeError as e:
+        with open(pickle_file, 'rb') as f:
+            pickle_data = pickle.load(f, encoding='latin1')
+    except Exception as e:
+        print('Unable to load data ', pickle_file, ':', e)
+        raise
+    return pickle_data
+
+def sym_adj(adj):
+    """Symmetrically normalize adjacency matrix."""
+    adj = sp.coo_matrix(adj)
+    rowsum = np.array(adj.sum(1))
+    d_inv_sqrt = np.power(rowsum, -0.5).flatten()
+    d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.
+    d_mat_inv_sqrt = sp.diags(d_inv_sqrt)
+    return adj.dot(d_mat_inv_sqrt).transpose().dot(d_mat_inv_sqrt).astype(np.float32).todense()
+
+def asym_adj(adj):
+    adj = sp.coo_matrix(adj)
+    rowsum = np.array(adj.sum(1)).flatten()
+    d_inv = np.power(rowsum, -1).flatten()
+    d_inv[np.isinf(d_inv)] = 0.
+    d_mat = sp.diags(d_inv)
+    return d_mat.dot(adj).astype(np.float32).todense()
+
+def calculate_normalized_laplacian(adj):
+    """
+    # L = D^-1/2 (D-A) D^-1/2 = I - D^-1/2 A D^-1/2
+    # D = diag(A 1)
+    :param adj:
+    :return:
+    """
+    adj = sp.coo_matrix(adj)
+    d = np.array(adj.sum(1))
+    d_inv_sqrt = np.power(d, -0.5).flatten()
+    d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.
+    d_mat_inv_sqrt = sp.diags(d_inv_sqrt)
+    normalized_laplacian = sp.eye(adj.shape[0]) - adj.dot(d_mat_inv_sqrt).transpose().dot(d_mat_inv_sqrt).tocoo()
+    return normalized_laplacian
+
+def calculate_random_walk_matrix(adj_mx):
+    adj_mx = sp.coo_matrix(adj_mx)
+    d = np.array(adj_mx.sum(1))
+    d_inv = np.power(d, -1).flatten()
+    d_inv[np.isinf(d_inv)] = 0.
+    d_mat_inv = sp.diags(d_inv)
+    random_walk_mx = d_mat_inv.dot(adj_mx).tocoo()
+    return random_walk_mx
+
+def calculate_reverse_random_walk_matrix(adj_mx):
+    return calculate_random_walk_matrix(np.transpose(adj_mx))
+
+def calculate_scaled_laplacian(adj_mx, lambda_max=2, undirected=True):
+    if undirected:
+        adj_mx = np.maximum.reduce([adj_mx, adj_mx.T])
+    L = calculate_normalized_laplacian(adj_mx)
+    if lambda_max is None:
+        lambda_max, _ = linalg.eigsh(L, 1, which='LM')
+        lambda_max = lambda_max[0]
+    L = sp.csr_matrix(L)
+    M, _ = L.shape
+    I = sp.identity(M, format='csr', dtype=L.dtype)
+    L = (2 / lambda_max * L) - I
+    return L.astype(np.float32)
+
+def load_adj(pkl_filename, adjtype):
+    if "PEMS0" in pkl_filename or "D7" in pkl_filename:
+        adj_mx = load_pickle(pkl_filename)
+    else:
+        sensor_ids, sensor_id_to_ind, adj_mx = load_pickle(pkl_filename)
+    if adjtype == "scalap":
+        adj = [calculate_scaled_laplacian(adj_mx)]
+    elif adjtype == "normlap":
+        adj = [calculate_normalized_laplacian(adj_mx).astype(np.float32).todense()]
+    elif adjtype == "symadj":
+        adj = [sym_adj(adj_mx)]
+    elif adjtype == "transition":
+        adj = [asym_adj(adj_mx)]
+    elif adjtype == "doubletransition":
+        adj = [asym_adj(adj_mx), asym_adj(np.transpose(adj_mx))]
+    elif adjtype == "identity":
+        adj = [np.diag(np.ones(adj_mx.shape[0])).astype(np.float32)]
+    else:
+        error = 0
+        assert error, "adj type not defined"
+    return adj
+
+def print_params(model):
+    # print trainable params
+    param_count = 0
+    print('Trainable parameter list:')
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            print(name, param.shape, param.numel())
+            param_count += param.numel()
+    print(f'\n In total: {param_count} trainable parameters. \n')
+    return