<jupyter_start><jupyter_text><jupyter_code>import os, sys
#from google.colab import drive
#drive.mount('/content/drive')
#nb_path = '/content/libraries'
sys.path.append('/content/drive/My Drive/Colab Notebooks/VGRNN/')
sys.path.append('/content/drive/My Drive/Colab Notebooks/')

#os.symlink('/content/drive/My Drive/Colab Notebooks', nb_path)
#sys.path.insert(0,nb_path)
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

#!python --version   #3.7.10

import io
import math
import numpy as np
import torch
import torch.nn as nn
import torch.utils
import torch.utils.data
from torchvision import datasets, transforms
from torch.autograd import Variable
import matplotlib.pyplot as plt 
from scipy.ndimage import rotate
from torch.distributions.uniform import Uniform
from torch.distributions.normal import Normal
#from sklearn.datasets import fetch_mldata
# from torch_geometric import nn as tgnn
from input_data import load_data
from preprocessing import preprocess_graph, construct_feed_dict, sparse_to_tuple, mask_test_edges
import scipy.sparse as sp
from scipy.linalg import block_diag
from torch.nn.parameter import Parameter
from torch.nn.modules.module import Module
import tarfile
import torch.nn.functional as F
import copy
import time

#print(torch.__version__)        #1.9.0+cu102
#print(torch.version.cuda)       #10.2

#!pip uninstall torch-scatter torch-sparse torch-geometric
!pip install -q torch-scatter -f https://pytorch-geometric.com/whl/torch-1.6.0+cu102.html
!pip install -q torch-sparse -f https://pytorch-geometric.com/whl/torch-1.6.0+cu102.html
!pip install -q torch-geometric

import torch_scatter
from torch_scatter import scatter_mean, scatter_max, scatter_add
from torch_geometric.utils import remove_self_loops, add_self_loops, degree
#from torch_geometric.datasets import Planetoid
import networkx as nx
import scipy.io as sio

import inspect
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score, roc_auc_score, average_precision_score
from sklearn.manifold import TSNE

import copy
import pickle
!pip install sparse
import sparse

import time
import datetime
from datetime import timedelta
import pandas as pd

from torch.utils.tensorboard import SummaryWriter

from warnings import simplefilter
# ignore all future warnings
simplefilter(action='ignore', category=FutureWarning)


seed = 3
np.random.seed(seed)


# utility functions

def uniform(size, tensor):
    stdv = 1.0 / math.sqrt(size)
    if tensor is not None:
        tensor.data.uniform_(-stdv, stdv)


def glorot(tensor):
    stdv = math.sqrt(6.0 / (tensor.size(0) + tensor.size(1)))
    if tensor is not None:
        tensor.data.uniform_(-stdv, stdv)


def zeros(tensor):
    if tensor is not None:
        tensor.data.fill_(0)


def ones(tensor):
    if tensor is not None:
        tensor.data.fill_(1)


def reset(nn):
    def _reset(item):
        if hasattr(item, 'reset_parameters'):
            item.reset_parameters()

    if nn is not None:
        if hasattr(nn, 'children') and len(list(nn.children())) > 0:
            for item in nn.children():
                _reset(item)
        else:
            _reset(nn)

def tuple_to_array(lot):
    out = np.array(list(lot[0]))
    for i in range(1, len(lot)):
        out = np.vstack((out, np.array(list(lot[i]))))
    
    return out


# masking functions

def mask_edges_det(adjs_list):

    adj_train_l, train_edges_l, val_edges_l = [], [], []
    val_edges_false_l, test_edges_l, test_edges_false_l = [], [], []
    edges_list = []
    for i in range(0, len(adjs_list)):
        # Function to build test set with 10% positive links
        # NOTE: Splits are randomized and results might slightly deviate from reported numbers in the paper.
        
        adj = adjs_list[i]
        # Remove diagonal elements
        adj = adj - sp.dia_matrix((adj.diagonal()[np.newaxis, :], [0]), shape=adj.shape)
        adj.eliminate_zeros()

        # Check that diag is zero:
        assert np.diag(adj.todense()).sum() == 0
        
        #get the upper trianglar portion of the matrix.
        adj_triu = sp.triu(adj)

        #convert the matrix into a tuple of the format: ((([1, 10]), ([1, 1, 1,..., 1, 1, 1])),...)
        adj_tuple = sparse_to_tuple(adj_triu)

        #get only the 0 index of the tuple. Returns as list: [[1 10],[1 12],[1 4],[20 25]]
        #shape: (n, 2)
        edges = adj_tuple[0]

        #convert the adj sparse matrix to tuple and return the result of the 0 index of the tuple
        edges_all = sparse_to_tuple(adj)[0]

        #get the number of test set: row number(n)/10
        num_test = int(np.floor(edges.shape[0] / 10.))
      
        #get the number of the validation set: row number(n)/20
        num_val = int(np.floor(edges.shape[0] / 20.))
    
        #list numbers of edge index based on the row axis of the edges
        #all_edge_idx = range(edges.shape[0])
        all_edge_idx = list(range(edges.shape[0]))

        #randomize the result
        np.random.shuffle(all_edge_idx)

        #get validation edge index from the randomized edge list. Extract only numbers equal to num_val
        val_edge_idx = all_edge_idx[:num_val]

        #get test edge index from the randomized edge list.
        #Extract only numbers equal to [num_val : (num_val + num_test)]
        test_edge_idx = all_edge_idx[num_val:(num_val + num_test)]

        #get the main test edge set by extracting values fom the edge list indexed by the test_edge_idx list
        test_edges = edges[test_edge_idx]
        
        #get the main validation edge set by extracting values fom the edge list indexed by the test_edge_idx list
        val_edges = edges[val_edge_idx]

        #delete the stacked test and validation edge set (along the axis=0) from the list of edges. 
        #This will be the training set
        # [[162 165], [162 169], [162 172], [171 174]]

        train_edges = np.delete(edges, np.hstack([test_edge_idx, val_edge_idx]), axis=0)
        
        #append the list of main edges
        edges_list.append(edges)
        
        def ismember(a, b, tol=5):
            #Test whether all array elements along a given axis evaluate to True. (np.all)
            rows_close = np.all(np.round(a - b[:, None], tol) == 0, axis=-1)
            return np.any(rows_close)  #np.any evaluate whether any elements evaluate to True

        #get false edge test set
        test_edges_false = []
        #Do while test_egde_false list length is still less than the tst_edge list
        while len(test_edges_false) < len(test_edges):
            #get random integers between 0 (lower) and the row size of the adj (higher)
            idx_i = np.random.randint(0, adj.shape[0])
            idx_j = np.random.randint(0, adj.shape[0])

            #if right and left values are equal, go back to the top loop
            if idx_i == idx_j:
                continue
            #if the tuple of the 2 values are part of edges_all (returns a bool), back to top
            if ismember([idx_i, idx_j], edges_all):
                continue
            #if the empty test_edges_false list is not None, check the conditions
            if test_edges_false:
                #if the tuple of the 2 values are part of test_edges_false list, back to top
                if ismember([idx_j, idx_i], np.array(test_edges_false)):
                    continue
                if ismember([idx_i, idx_j], np.array(test_edges_false)):
                    continue
            #append result to the test_edges_false list
            test_edges_false.append([idx_i, idx_j])  #result sample: [[19, 2], [177, 163], [15, 119], [3, 155],...] 

        
        #get false validation edge set    
        val_edges_false = []
        while len(val_edges_false) < len(val_edges):
            idx_i = np.random.randint(0, adj.shape[0])
            idx_j = np.random.randint(0, adj.shape[0])
            if idx_i == idx_j:
                continue
            if ismember([idx_i, idx_j], train_edges):
                continue
            if ismember([idx_j, idx_i], train_edges):
                continue
            if ismember([idx_i, idx_j], val_edges):
                continue
            if ismember([idx_j, idx_i], val_edges):
                continue
            if val_edges_false:
                if ismember([idx_j, idx_i], np.array(val_edges_false)):
                    continue
                if ismember([idx_i, idx_j], np.array(val_edges_false)):
                    continue
            val_edges_false.append([idx_i, idx_j])

        r""" The assert keyword lets you test if a condition in your code returns True, 
        if not, the program will raise an AssertionError.

        #we assert the truthfulness of these conditions. 
        #check to confirm that the values (arg: 1) are bitwise NOT (tilde)
        #in the set of values (arg: 2) in the other list."""

        assert ~ismember(test_edges_false, edges_all)
        assert ~ismember(val_edges_false, edges_all)
        assert ~ismember(val_edges, train_edges)
        assert ~ismember(test_edges, train_edges)
        assert ~ismember(val_edges, test_edges)

        #get np.ones of elements of the row size of the train_edges
        data = np.ones(train_edges.shape[0])

        # Re-build adj matrix for the training set
        r""" [ : , 0 ] means (more or less) [ first_row:last_row , column_0 ]. 
        If you have a 2-dimensional list/matrix/array, this notation will give you all 
        the values in column 0 (from all rows)."""

        adj_train = sp.csr_matrix((data, (train_edges[:, 0], train_edges[:, 1])), shape=adj.shape)

        #add the new adjacency matrix to its transpose
        adj_train = adj_train + adj_train.T

        #fill all the initialised list
        adj_train_l.append(adj_train)
        train_edges_l.append(train_edges)
        val_edges_l.append(val_edges)
        test_edges_l.append(test_edges)
        val_edges_false_l.append(val_edges_false)
        test_edges_false_l.append(test_edges_false)

    # NOTE: these edge lists only contain single direction of edge!
    return adj_train_l, train_edges_l, val_edges_l, val_edges_false_l, test_edges_l, test_edges_false_l

def mask_edges_prd(adjs_list):
    pos_edges_l , false_edges_l = [], []
    edges_list = []
    for i in range(0, len(adjs_list)):
        # Function to build test set with 10% positive links
        # NOTE: Splits are randomized and results might slightly deviate from reported numbers in the paper.
        
        adj = adjs_list[i]
        # Remove diagonal elements
        adj = adj - sp.dia_matrix((adj.diagonal()[np.newaxis, :], [0]), shape=adj.shape)
        adj.eliminate_zeros()
        # Check that diag is zero:
        assert np.diag(adj.todense()).sum() == 0
        
        adj_triu = sp.triu(adj)
        adj_tuple = sparse_to_tuple(adj_triu)
        edges = adj_tuple[0]
        edges_all = sparse_to_tuple(adj)[0]
        num_false = int(edges.shape[0])
        
        pos_edges_l.append(edges)
        
        def ismember(a, b, tol=5):
            rows_close = np.all(np.round(a - b[:, None], tol) == 0, axis=-1)
            return np.any(rows_close)
        
        edges_false = []
        while len(edges_false) < num_false:
            idx_i = np.random.randint(0, adj.shape[0])
            idx_j = np.random.randint(0, adj.shape[0])
            if idx_i == idx_j:
                continue
            if ismember([idx_i, idx_j], edges_all):
                continue
            if edges_false:
                if ismember([idx_j, idx_i], np.array(edges_false)):
                    continue
                if ismember([idx_i, idx_j], np.array(edges_false)):
                    continue
            edges_false.append([idx_i, idx_j])

        assert ~ismember(edges_false, edges_all)
        
        false_edges_l.append(edges_false)

    # NOTE: these edge lists only contain single direction of edge!
    return pos_edges_l, false_edges_l    
# loading data

path = 'drive/My Drive/Colab Notebooks/VGRNN/data/'
# # Enron dataset
with open(path+'enron_data/enron_adj_sparse_matrix.pickle', 'rb') as handle:
     adj_sparse_matrix = pickle.load(handle)

with open(path+'enron_data/enron_adj_dense_matrix.pickle', 'rb') as handle:
     adj_dense_matrix = pickle.load(handle)

with open(path+'enron_data/enron_edge_attribute_matrix.pickle', 'rb') as handle:
     edge_attr_matrix = pickle.load(handle)

with open(path+'enron_data/enron_node_attribute_matrix.pickle', 'rb') as handle:
     node_attr_matrix = pickle.load(handle)

adj_sparse_matrix = adj_sparse_matrix[7:34]                   
adj_dense_matrix = adj_dense_matrix[7:34]         
edge_attr_matrix = edge_attr_matrix[7:34]          
node_attr_matrix = node_attr_matrix[7:34]           


outs = mask_edges_det(adj_sparse_matrix)

#reconstructed adjacency matrix of the training set
adj_train_l = outs[0]              

#List of training edge set
train_edges_l = outs[1]              

#List of validation edge set
val_edges_l = outs[2]               

#List of false validation edge set(i.e., never exist)
val_edges_false_l = outs[3]  
      

#List of test edge set
test_edges_l = outs[4]          

#List of false test edge set 
test_edges_false_l = outs[5]         


pos_edges_l, false_edges_l = mask_edges_prd(adj_sparse_matrix)
# creating edge list
edge_idx_list = []                 

for i in range(len(train_edges_l)):
    edge_idx_list.append(torch.tensor(np.transpose(train_edges_l[i]), dtype=torch.long))

#print('Training edges: ', edge_idx_list)
# layers

class E_GCN_Conv(nn.Module):
    def __init__(self, in_channels, out_channels, act=F.relu, improved=True, bias=True, num_channels=10, aggr='sum'):
        super(E_GCN_Conv, self).__init__()

        self.in_channels = in_channels            #[64]
        self.out_channels = out_channels          #[32]
        self.act = act
        self.num_channels = num_channels

        self.weight = Parameter(torch.Tensor(in_channels, out_channels, num_channels))

        if bias:
            self.bias = Parameter(torch.Tensor(out_channels, num_channels))
        else:
            self.register_parameter('bias', None)

        self.reset_parameters()

        if (aggr == 'concat'):
            self.aggr = 'concat'
            self.last_ops = nn.Linear(self.out_channels * self.num_channels, self.out_channels)
        elif (aggr == 'sum'):
            self.aggr = 'sum'
            self.last_ops = nn.Linear(self.out_channels, self.out_channels)
        

    def reset_parameters(self):
        glorot(self.weight)
        zeros(self.bias)

    
    def forward(self, x, edge_index, edge_attr):

        #add or remove node self loop. We remove in our case
        edge_index, edge_attr = remove_self_loops(edge_index, edge_attr)

        
        #edge index rows and column representation
        row, col = edge_index     #[21]

        #normalize the adjacency matrix
        #deg = scatter_add(edge_attr, row, dim=0, dim_size=x.size(0))
        deg = degree(col, x.size(0), dtype=x.dtype)
        deg_inv_sqrt = deg.pow(-0.5)
        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0

        #reshape the row and column vectors
        deg_inv_sqrt_row = deg_inv_sqrt[row].view(-1, 1)      #[[1.0000],[1.0000]]
        deg_inv_sqrt_col = deg_inv_sqrt[col].view(-1, 1)      #[[0.5774],[0.0000]]

        #multiply row and col vectors with edge weights (We replace the adjacency matrix with the edge tensor)
        norm_edge = deg_inv_sqrt_row * edge_attr * deg_inv_sqrt_col     #size([edge_index[row/col] No., 14])


        #Slice and list the normalized vectors based on the nu. of channels
        norm = []
        for i in range(0, edge_attr.size()[1]):
            norm.append(norm_edge[:, i:i+1])

        node_state_list = []
        #for each edge channels, we perform a weighted convolution with edge weights as co-efficient
        for c in range(self.num_channels):
            if self.in_channels > self.out_channels:

                #if the weight matrix is not none
                if self.weight is not None:
                    #matrix product of the node (hidden state) with the weight matrix
                    weighted_nodes = torch.matmul(x, self.weight[:, :, c])      #(size[149, 32])
                else:
                    #otherwise, hidden state remains same
                    weighted_nodes = x
                
                #if vectors are normalized
                if norm is not None: 
                    #multiply each element in the each channels of the norm with weighted hidden state             
                    weighted_conv = torch.mul(norm[c], weighted_nodes[row])      #size(21, 32)

                    #propagate messages through all edges and update the nodes
                    weighted_conv_sum = scatter_add(weighted_conv, col, dim=0, dim_size=x.size(0)) #size(149, 32)
                else:
                    weighted_conv_sum = scatter_add(weighted_nodes[row], col, dim=0, dim_size=x.size(0))

                channel_node_state = weighted_conv_sum

            else:
                if norm is not None:
                    unweighted_conv = torch.mul(norm[c], x[row])
                    unweighted_conv_sum = scatter_add(unweighted_conv, col, dim=0, dim_size=x.size(0))
                else:
                    unweighted_conv_sum = scatter_add(x[row], col, dim=0, dim_size=x.size(0))
                
                if self.weight is not None:
                    channel_node_state = torch.matmul(unweighted_conv_sum.float(), self.weight[:, :, c])
            
            #add linear bias if True
            if self.bias is not None:
                channel_node_state = channel_node_state + self.bias[:, c]
            
            #pass param through a linear activation function
            channel_node_state = self.act(channel_node_state)
            #append each channel to node state list
            node_state_list.append(channel_node_state)        #size(N, 32/16)

        #we consider two aggregation method across each channels of the edge weights
        #1. Sum aggregation method 
        if (self.aggr == 'sum'):
            node_states = torch.stack(node_state_list, dim=1).sum(1).float()     #[N, 32]
        #2. Concat aggregation method               
        elif (self.aggr == 'concat'): 
            node_states = torch.cat(node_state_list, dim=1).float()
            
        #pass aggregated vectors through a flexible linear transformation layer      
        out = self.last_ops(node_states)                        #size(N, 32/16)

        return out

    def __repr__(self):
        return '{}({}, {})'.format(self.__class__.__name__, self.in_channels, 
                                   self.out_channels, self.num_channels)
class gru_gcn(nn.Module):
    def __init__(self, input_size, hidden_size, n_layer, bias=True):
        super(gru_gcn, self).__init__()

        self.hidden_size = hidden_size
        self.n_layer = n_layer
        
        # gru weights
        self.weight_xz = []
        self.weight_hz = []
        self.weight_xr = []
        self.weight_hr = []
        self.weight_xh = []
        self.weight_hh = []
        
        for i in range(self.n_layer):
            if i==0:
                self.weight_xz.append(E_GCN_Conv(input_size, hidden_size, act=lambda x:x, bias=bias))
                self.weight_hz.append(E_GCN_Conv(hidden_size, hidden_size, act=lambda x:x, bias=bias))
                self.weight_xr.append(E_GCN_Conv(input_size, hidden_size, act=lambda x:x, bias=bias))
                self.weight_hr.append(E_GCN_Conv(hidden_size, hidden_size, act=lambda x:x, bias=bias))
                self.weight_xh.append(E_GCN_Conv(input_size, hidden_size, act=lambda x:x, bias=bias))
                self.weight_hh.append(E_GCN_Conv(hidden_size, hidden_size, act=lambda x:x, bias=bias))
            else:
                self.weight_xz.append(E_GCN_Conv(hidden_size, hidden_size, act=lambda x:x, bias=bias))
                self.weight_hz.append(E_GCN_Conv(hidden_size, hidden_size, act=lambda x:x, bias=bias))
                self.weight_xr.append(E_GCN_Conv(hidden_size, hidden_size, act=lambda x:x, bias=bias))
                self.weight_hr.append(E_GCN_Conv(hidden_size, hidden_size, act=lambda x:x, bias=bias))
                self.weight_xh.append(E_GCN_Conv(hidden_size, hidden_size, act=lambda x:x, bias=bias))
                self.weight_hh.append(E_GCN_Conv(hidden_size, hidden_size, act=lambda x:x, bias=bias))
    
    def forward(self, inp, edge_index, edge_tensor, h):
        h_out = torch.zeros(h.size())
        for i in range(self.n_layer):
            if i==0:
                
                z_g = torch.sigmoid(self.weight_xz[i](inp, edge_index, edge_tensor) + self.weight_hz[i](h[i], edge_index, edge_tensor))
                r_g = torch.sigmoid(self.weight_xr[i](inp, edge_index, edge_tensor) + self.weight_hr[i](h[i], edge_index, edge_tensor))
                h_tilde_g = torch.tanh(self.weight_xh[i](inp, edge_index, edge_tensor) + self.weight_hh[i](r_g * h[i], edge_index, edge_tensor))
                h_out[i] = z_g * h[i][0: inp.size(0)] + (1 - z_g) * h_tilde_g
        #         out = self.decoder(h_t.view(1,-1))
            else:
                z_g = torch.sigmoid(self.weight_xz[i](h_out[i-1], edge_index, edge_tensor) + self.weight_hz[i](h[i], edge_index, edge_tensor))
                r_g = torch.sigmoid(self.weight_xr[i](h_out[i-1], edge_index, edge_tensor) + self.weight_hr[i](h[i], edge_index, edge_tensor))
                h_tilde_g = torch.tanh(self.weight_xh[i](h_out[i-1], edge_index, edge_tensor) + self.weight_hh[i](r_g * h[i], edge_index, edge_tensor))
                h_out[i] = z_g * h[i] + (1 - z_g) * h_tilde_g
        #         out = self.decoder(h_t.view(1,-1))
        
        out = h_out
        return out, h_out

# VGRNN model

class VGAE_Edge(nn.Module):
    def __init__(self, node_feat_dim, hidden_dim, latent_var_dim, n_layers, edge_feat_dim, eps, conv='GCN', bias=False):
        super(VGAE_Edge, self).__init__()
        
        #input dimension
        self.node_feat_dim = node_feat_dim        
        self.eps = eps
        #hidden_layer dim.
        self.hidden_dim = hidden_dim        #32
        #latent variable dim.
        self.latent_var_dim = latent_var_dim        #10
        self.n_layers = n_layers   #1
        self.edge_feat_dim = edge_feat_dim
       
        if conv == 'GCN':
            #flexible sequential neural network linear transformations
            self.input_emb = nn.Sequential(nn.Linear(node_feat_dim, hidden_dim), nn.ReLU())
            self.output_emb = nn.Sequential(nn.Linear(latent_var_dim, hidden_dim))
            
            #encoder functions
            self.encoder = E_GCN_Conv(hidden_dim + hidden_dim, hidden_dim)            
            self.encoder_mu = E_GCN_Conv(hidden_dim, latent_var_dim, act=lambda x:x)
            self.encoder_sigma = E_GCN_Conv(hidden_dim, latent_var_dim, act=F.softplus)
            
            #linear linear transformation of the prior functions
            self.prior = nn.Sequential(nn.Linear(hidden_dim, hidden_dim), nn.ReLU())
            self.prior_mu = nn.Sequential(nn.Linear(hidden_dim, latent_var_dim))
            self.prior_sigma = nn.Sequential(nn.Linear(hidden_dim, latent_var_dim), nn.Softplus())
            
            #recurrent neural networks model function
            self.rnn = gru_gcn(hidden_dim + hidden_dim, hidden_dim, n_layers, bias)          
          

    def forward(self, x, edge_idx_list, edge_attr_matrix, hidden_in=None):

        #assert the length of edge matrix = elngth of the edge indices
        assert len(edge_attr_matrix) == len(edge_idx_list)

        #print(x.size())        #[26, 149, 6]

        #initialize params
        kld_loss = 0
        l2_loss = 0
        encoder_mu_list, encoder_sigma_list = [], []
        prior_mu_list, prior_sigma_list = [], []
        decoded_list, z_list = [], []
        kld_loss_list, l2_loss_list = [], []
        
        #hidden var will be none in the first set of operations
        if hidden_in is None:
            #so we create a matrix of zeros as initial representation
            h = torch.zeros(self.n_layers, x.size(1), self.hidden_dim)  #size([1, 149, 32])
        else:
            #hidden var here will be the recurrent vectors
            h = hidden_in


        for t in range(x.size(0)):  

            #linearly transform x features
            input_emb_t = self.input_emb(x[t].float())              #[149, 32]
            #edge indices at time t    
            edge_idx_list_t = edge_idx_list[t]

            #edge tensor matrix at time t => extract on the tensors associated with the edge indices at time t
            edge_tensor_t = (edge_attr_matrix[t][edge_idx_list_t[0], edge_idx_list_t[1]])#[:, 0:latent_var_dim]

            
            #encoder
            #encoders conditioned on priors so features of previous states can be 
            #recurrently modeled
            encoder_t = self.encoder(torch.cat([input_emb_t, h[-1]], 1), edge_idx_list_t, edge_tensor_t)    #[149, 32]
            #encoder mean
            encoder_mu_t = self.encoder_mu(encoder_t, edge_idx_list_t, edge_tensor_t)                   #[149, 16]
            #encoder standard deviation
            encoder_sigma_t = self.encoder_sigma(encoder_t, edge_idx_list_t, edge_tensor_t)                     #[149, 16]
            
            #prior
            prior_t = self.prior(h[-1])                               #[149, 32]
            prior_mu_t = self.prior_mu(prior_t)                       #[149, 16]
            prior_sigma_t = self.prior_sigma(prior_t)                 #[149, 16]
            
            #sampling and reparameterization
            z_t = self._reparameterized_sample(encoder_mu_t, encoder_sigma_t)  #[149, 16]
            #apply a fully connected layer to z_t
            output_emb_t = self.output_emb(z_t)                                  #[149, 32]

            #decoder function -> takes the linearly transformed latent variable and egde indices as args
            #decoder_t = self.dec(z_t, edge_idx_list_t)
            decoder_t = self.dec(output_emb_t, edge_idx_list_t)

            #recurrence
            _, h = self.rnn(torch.cat([input_emb_t, output_emb_t], 1), edge_idx_list_t, edge_tensor_t, h)       #[1, 149, 32]
            
            num_nodes = edge_attr_matrix[t].size(0)
            encoder_mu_t_slice = encoder_mu_t[0:num_nodes, :]
            encoder_sigma_t_slice = encoder_sigma_t[0:num_nodes, :]
            prior_mu_t_slice = prior_mu_t[0:num_nodes, 0:num_nodes]
            prior_sigma_t_slice = prior_sigma_t[0:num_nodes, :]

            #computing losses
            kld_loss_t = self.kl_divergence(encoder_mu_t_slice, encoder_sigma_t_slice, prior_mu_t_slice, prior_sigma_t_slice)
            kld_loss_list.append(kld_loss_t)
            kld_loss = kld_loss + kld_loss_t
            #kld_loss += self.kl_divergence_zu(encoder_mu_t, encoder_sigma_t)
            l2_loss_t = self._l2_norm(decoder_t, edge_tensor_t)
            l2_loss_list.append(l2_loss_t)
            l2_loss = l2_loss + l2_loss_t
            
            encoder_sigma_list.append(encoder_sigma_t_slice)
            encoder_mu_list.append(encoder_mu_t_slice)
            prior_mu_list.append(prior_mu_t_slice)
            prior_sigma_list.append(prior_sigma_t_slice)
            decoded_list.append(decoder_t)
            z_list.append(z_t)

        #print(decoded_list)
        return kld_loss, l2_loss, encoder_mu_list, prior_mu_list, decoded_list, h, kld_loss_list, l2_loss_list
    
    #decoder function
    def dec(self, z, edge_index):
        #input features have dimension = the col sizes of Zi and Zj 
        in_feat = int(z.size(1))
        #output feature has the size of the edge channels
        out_feat = int(self.edge_feat_dim)
        #output = neural network decoder
        outputs = Decoder(in_feat, out_feat, act=lambda x:x)(z, edge_index)
        return outputs
    
    def reset_parameters(self, stdv=1e-1):
        for weight in self.parameters():
            weight.data.normal_(0, stdv)
     
    def _init_weights(self, stdv):
        pass
    
    def _reparameterized_sample(self, mean, std):
        eps1 = torch.FloatTensor(std.size()).normal_()
        eps1 = Variable(eps1)
        return eps1.mul(std).add_(mean)
    
    #VAE loss function regularizer
    def kl_divergence(self, encoder_mu, encoder_sigma, prior_mu, prior_sigma):
        mu_size = encoder_mu.size(0)
        encoder_sigma_log = torch.log(encoder_sigma + self.eps)
        prior_sigma_log = torch.log(prior_sigma + self.eps)
        encoder_sigma = encoder_sigma + self.eps
        prior_sigma = prior_sigma + self.eps

        kld_element = (2 * prior_sigma_log - 2 * encoder_sigma_log + (torch.pow(encoder_sigma, 2) + torch.pow(encoder_mu - prior_mu, 2)) / 
                      torch.pow(prior_sigma, 2) - 1)
        kld_element = kld_element.detach().numpy()      
        kld_element = torch.tensor(np.nan_to_num(kld_element, copy=True, nan=0.0))
        kld = (0.5 / mu_size) * kld_element.sum(1).mean()
        return kld
    
    def kl_divergence_zu(self, mu, sigma):        
        mu_size = mu.size(0)
        sigma_log = torch.log(sigma + self.eps)
        kld_element =  (1 + 2*sigma_log - (sigma**2) - (mu**2))
        kld_element = kld_element.detach().numpy()      
        kld_element = torch.tensor(np.nan_to_num(kld_element, copy=True, nan=0.0))        
        return (-0.5 / mu_size) * kld_element.sum(1).mean()
    
    def regularizer(self, samples_size, features_size, lambda_value=0.01):

        m = samples_size
        n = features_size
        W = torch.randn(m, n)

        reg_term = (lambda_value / (2 * m)) * torch.sum(torch.square(W))
        #print(reg_term)
        return reg_term

    #reconstruction loss - (nll for binary classification model)
    #the log likelihood of the true observation given the predicted distribution
    def _l2_norm(self, pred, actual):
        x_size_row = actual.size(0)
        x_size_col = actual.size(1)

        #l2_reg = self.regularizer(x_size_row, x_size_col)
        loss = nn.MSELoss(reduction='mean')
        l2_loss = loss(input=pred.float(), target=actual.float())
        l2_loss_val = (1.0 / x_size_row) * l2_loss
        l2_norm = l2_loss_val #+ l2_reg

        return l2_norm


class Decoder(nn.Module):
    def __init__(self, in_feat, out_feat, act=torch.sigmoid):
        super(Decoder, self).__init__()        

        self.act = act
        self.in_feat = in_feat
        self.out_feat = out_feat

        self.edge_nn_mean = nn.Sequential(nn.Linear(self.in_feat, self.out_feat))
        self.edge_nn_std = nn.Sequential(nn.Linear(self.in_feat, self.out_feat), nn.Softplus())

        self.edge_nn = nn.Linear(self.in_feat * 2, self.out_feat, bias=False)

    def forward(self, z, edge_index):
        z = F.dropout(z, p=0., training=True) 

        #emb_mean = self.edge_nn_mean(z)
        #emb_std = self.edge_nn_std(z)

        #x_hat  = emb_std.add(emb_mean)         #size[149, 10]

        #embeddings of edge_index 0
        z0_emb = z[edge_index[0]]       #size[N, 10]
        #embeddings of edge index 1     
        z1_emb = z[edge_index[1]]

        #concatenate the embeddings of the edge indices
        #edge_mult = (z0_emb * z1_emb)#.sum(dim=1)      #size[N, 20]

        edge_mult = torch.cat([z0_emb, z1_emb], dim=-1)

        r'''pass through a neural network. Sigmoid activation function can be used 
          in case of binomial cross-entropy problem.'''

        #For regression task, just a linear or identity function is okay
        edge_emb = self.edge_nn(edge_mult)                   #size[N, 10]

        return edge_emb
r""" Calculate and evaluate the Area Under (Receiver Operating Characteristic) Curve 
     and the Average Precision (AP) """

# evaluation function

def get_edge_eval_scores(edges_pos, edges_neg, edge_attr_matrix, edge_feat_dim, edge_emb):
    
    mse, rmse = [], []
    mae, mape = [], []


    in_channel = edge_emb[0].size(1)
    out_channel = edge_feat_dim 

    nn_layer = nn.Linear(in_channel*2, out_channel, bias=False)

    for i in range(len(edges_pos)):
        
        # Predict on test set of edges
        #explicitly remove the computational graph of the tensor 
        #(from gradient descent) with detach and change back to numpy
        emb = edge_emb[i].detach().numpy()

        z_emb = torch.tensor(emb)

        #edge tensor matrix
        edge_attr_mat = edge_attr_matrix[i]
       
        #initialize predicted edge list
        pred_pos_edges, pos_edges = [], []
        pos_edge_list = []
        pos_list = []

        for e in edges_pos[i]:
  
            z_i = z_emb[e[0]]
            z_j = z_emb[e[1]]

            cat_embs = torch.cat([z_i, z_j], dim=-1) 
            pos_embs_var = nn_layer(cat_embs).detach().numpy()
            #append the sigmoid computation of the reconstructed embedding matrix
            #Note: we can also consider without sigmoid as our task is regression like
            pred_pos_edges.append(pos_embs_var)
            #positive edge tensor
            pos_edges.append((edge_attr_mat[e[0], e[1]]))
            pos_edge_list.append((e[0], e[1], pos_embs_var))
            pos_list.append((e[0], e[1], (edge_attr_mat[e[0], e[1]])))
 
        pred_neg_edges, neg_edges = [], []
        neg_edge_list = []
        neg_list = []

        for e in edges_neg[i]:

            z_i = z_emb[e[0]]
            z_j = z_emb[e[1]]

            cat_embs = torch.cat([z_i, z_j], dim=-1) 
            neg_embs_var = nn_layer(cat_embs).detach().numpy()
            pred_neg_edges.append(neg_embs_var)  

            neg_edges.append((edge_attr_mat[e[0], e[1]]))
            neg_edge_list.append((e[0], e[1], neg_embs_var))
            neg_list.append((e[0], e[1], (edge_attr_mat[e[0], e[1]])))

        #stack up the positive and negative predicted features
        pred_all_edges = np.hstack([pred_pos_edges, pred_neg_edges])    

        #for error free mean square eval
        pos_edges = [t.detach().numpy() for t in pos_edges]
        neg_edges = [t.detach().numpy() for t in neg_edges]

        #stack up all positive and negative features
        all_true_edges = np.hstack([pos_edges, neg_edges])

        #evaluate the mean square error loss of the ground truth and predicted values
        mse.append(mean_squared_error(all_true_edges, pred_all_edges))
        rmse.append(mean_squared_error(all_true_edges, pred_all_edges, squared=False))
        mae.append(mean_absolute_error(all_true_edges, pred_all_edges))
        mape.append(mean_absolute_error(all_true_edges, pred_all_edges)*100)

    return mse, rmse, mae, mape, pos_edge_list, neg_edge_list, pos_list, neg_list
# hyperparameters

hidden_dim = 32
latent_var_dim = 16
n_layers =  1
clip = 10
learning_rate = 1e-2
num_nodes = node_attr_matrix[0].shape[1]
node_feat_dim = num_nodes
edge_feat_dim = 10
timesteps_len = len(train_edges_l)  #27
#print(timesteps_len)
eps = 1e-10
conv_type='GCN'

# creating input tensors
node_attr = torch.stack(node_attr_matrix)   #[80, 149, 6]


adj_label_list = []
for i in range(len(adj_train_l)):
    temp_matrix = adj_train_l[i]
    adj_label_list.append(torch.tensor(temp_matrix.toarray().astype(np.float32))) 

# building model
model = VGAE_Edge(node_feat_dim, hidden_dim, latent_var_dim, n_layers, edge_feat_dim, eps, conv=conv_type, bias=True)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

#print(model)

# training
timesteps_init = 0
#25% of 36
timesteps_end = timesteps_len - 1      #27
#print(timesteps_end)
test_init = 0

#train_edges_l
training_edges = (train_edges_l[timesteps_end: timesteps_len])

edge_train = (edge_attr_matrix[timesteps_end: timesteps_len])



#writer = SummaryWriter('drive/MyDrive/Colab Notebooks/VGRNN/tensorboard_log/' + datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))

start_time = time.monotonic()
for k in range(1000):
    optimizer.zero_grad()

    kld_loss, l2_loss, _, _, _, hidden_state, kld_list, loss_list = model(node_attr[timesteps_init:timesteps_end]
                                                , edge_idx_list[timesteps_init:timesteps_end]
                                                , edge_attr_matrix[timesteps_init:timesteps_end]
                                                )
    
    #print(kld_list)
    #print(loss_list)

    loss = kld_loss + l2_loss
    loss.backward()
    optimizer.step()
       
    nn.utils.clip_grad_norm_(model.parameters(), clip)
    
    if k > test_init:
        kld_loss, l2_loss, encs_, priors_, edge_dec, _, kld_list, loss_list = model(node_attr[timesteps_end:timesteps_len]
                                              , edge_idx_list[timesteps_end:timesteps_len]
                                              , edge_attr_matrix[timesteps_end:timesteps_len]
                                              , hidden_state)

        mse_val, rmse_val, mae_val, mape_val, pred_pos, pred_neg, pos, neg = get_edge_eval_scores(
                                                          pos_edges_l[timesteps_end:timesteps_len]
                                                        , false_edges_l[timesteps_end:timesteps_len]
                                                        , edge_attr_matrix[timesteps_end:timesteps_len]
                                                        , edge_feat_dim
                                                        , priors_
                                                        )


        #mse_val, rmse_val, mae_val, mape_val, pred_pos, pred_neg, pos, neg = get_edge_eval_scores(
                                                         # val_edges_l[timesteps_end:timesteps_len]
                                                        #, val_edges_false_l[timesteps_end:timesteps_len]
                                                        #, edge_attr_matrix[timesteps_end:timesteps_len]
                                                        #, edge_feat_dim
                                                        #, priors_
                                                        #)

        #mse_test, rmse_test, mae_test, mape_test, pred_pos, pred_neg, pos, neg = get_edge_eval_scores(
                                                        #  test_edges_l[timesteps_end:timesteps_len]
                                                        #, test_edges_false_l[timesteps_end:timesteps_len]
                                                        #, edge_attr_matrix[timesteps_end:timesteps_len]
                                                        #, edge_feat_dim
                                                        #, priors_
                                                        #)        


    #Note: Prior mean reduces the loss than the decoded variables. 
    print('********************************************************')
    print('epoch: ', k)
    print('\nLOSS => kld_loss: {} | l2_loss: {} | loss: {}'.format(round(kld_loss.mean().item(), 4)
                                                                    , round(l2_loss.mean().item(), 4)
                                                                    , round(loss.mean().item(), 4)
                                                                    ))
    #writer.add_scalar("Loss/train", loss.mean().item(), k)
    if k > test_init:
        #writer.add_scalar("validation mean_score", np.mean(np.array(mse_val)), k)
        #writer.add_scalar("test mean_score", np.mean(np.array(mse_test)), k)

        
        print('\nEDGE RECONSTRUCTION VAL => mse: {} | rmse: {} | mae: {} | mape: {}'.format(
                                                                            round(np.mean(np.array(mse_val)), 4)
                                                                          , round(np.mean(np.array(rmse_val)), 4)
                                                                          , round(np.mean(np.array(mae_val)), 4)
                                                                          , round(np.mean(np.array(mape_val)), 4)
                                                                          ))
        #print('\nEDGE RECONSTRUCTION TEST => mse: {} | rmse: {} | mae: {} | mape: {}'.format(
                                                                          #  round(np.mean(np.array(mse_test)), 4)
                                                                          #, round(np.mean(np.array(rmse_test)), 4)
                                                                          #, round(np.mean(np.array(mae_test)), 4)
                                                                          #, round(np.mean(np.array(mape_test)), 4)
                                                                          #))


#print('Pos: ', pos)                                                                     
#print('Pos_Pred: ', pred_pos) 
#print('\nNeg: ', neg)
#print('Neg Pred: ', pred_neg)   
#writer.flush() 
#writer.close()   
end_time = time.monotonic()
print('Total Execution Time: {}'.format(timedelta(seconds=end_time - start_time)))

#!pip install tensorboard
#!tensorboard --logdir='drive/MyDrive/Colab Notebooks/VGRNN/tensorboard_log/'<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09974387474025033, 'fraction_numerical': 0.012709611946068719, 'mean_word_length': 2.1176647834274953, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 14, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 90, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '63163', 'license': 'MIT', 'n_tokens_mistral': 13374, 'n_tokens_neox': 12730, 'n_words': 2948, 'path': 'edge_reconstruction_main.ipynb', 'repo_name': 'spyderweb-abdul/Pattern-Recognition-and-Reconstruction-Detecting-Malicious-Deletions-in-Textual-Communications', 'size': '41386'}