init

HeteroVG_MNIST.py

@@ -0,0 +1,211 @@
+import os
+import re
+import json
+import numpy as np
+import pandas as pd
+
+import torch
+import random
+import pickle as pkl
+from tqdm import tqdm
+from torch import Tensor
+from scipy.spatial import distance_matrix
+import torch_geometric
+from torch_geometric.data import HeteroData
+from torch_geometric.nn import to_hetero
+# from shapely.geometry import Point, Polygon
+
+
+def cross_product(p1, p2, p3):
+    return (p2[0] - p1[0]) * (p3[1] - p1[1]) - (p3[0] - p1[0]) * (p2[1] - p1[1])
+
+def colinear(p1, p2, p3):
+  if (p1[1]-p2[1])*(p2[0]-p3[0]) == (p1[0]-p2[0])*(p2[1]-p3[1]) and p3[0]>min(p1[0],p2[0]) and p3[0]<max(p1[0],p2[0]): return True
+  if (p1[1]-p2[1])*(p2[0]-p3[0]) == (p1[0]-p2[0])*(p2[1]-p3[1]) and p3[1]>min(p1[1],p2[1]) and p3[1]<max(p1[1],p2[1]): return True
+
+def is_intersected(p1, p2, p3, p4):
+  if colinear(p1, p2, p3) or colinear(p1, p2, p4): return True
+  if cross_product(p1, p2, p3) * cross_product(p1, p2, p4) < 0 and cross_product(p3, p4, p1) * cross_product(p3, p4, p2) < 0:  return True
+  else: return False
+
+#Pos of single number
+def read_singe(df, i):
+  p_i = df[0][i]
+  np_i = list(p_i)
+  rflag = 0
+  for x in range(len(p_i)-4):
+    if p_i[x] == ')': rflag+=1
+    if p_i[x:x+4] == '), (' and p_i[x-1]!=')': np_i.insert(x+rflag+1,' 0.0 0.0')
+    elif p_i[x:x+4] == '), (' and p_i[x-1]==')': np_i.insert(x+rflag+1,' 1.0 1.0')
+  p_i = ''.join(np_i)
+
+
+  pos = np.empty((1,2))
+  pi_nums = re.findall(r"\d+\.?\d*",p_i)
+  j=0
+  while j < len(pi_nums)-1:
+    if j == 0:
+      pos[0][0] = float(pi_nums[j])
+      pos[0][1] = float(pi_nums[j+1])
+      j+=2
+      continue
+    pos = np.append(pos,[[float(pi_nums[j]),0]],0)
+    pos[j//2][1] = float(pi_nums[j+1])
+    j+=2
+  return pos
+
+def Visi_Edge(pos_join, flag):
+  inside_edge_index = [[],[]]
+  apart_edge_index = [[],[]]
+
+  vg_point = []
+  for i in range(len(pos_join)): vg_point.append((pos_join[i][0], pos_join[i][1]))
+
+  hole_p = np.where(flag==1)[0]
+  if len(hole_p) != 0:
+    last_id = 0
+    for m in range(len(flag)):
+      if flag[m] == 2 or flag[m] == 3:
+        if sum(flag[last_id:m]) == 0:
+          last_id = m+1
+          continue
+        poly_i = vg_point[last_id:m+1]
+        pos_i = np.arange(last_id, m+1)
+        last_id = m+1
+        for i in range(len(poly_i)):
+          if flag[pos_i[i]] == 1:
+            for j in range(i, len(flag)):
+              if flag[j]==1 or flag[j] == 2 or flag[j] == 3:
+                hole_i = poly_i[i:j+1]
+                pos_hole = np.arange(i, j+1)
+                for p1 in hole_i:
+                  for p2 in poly_i:
+                    if p2 not in hole_i:
+                      inter_count = 0
+                      for d in range(len(poly_i)-1):
+                        p3, p4 = poly_i[d], poly_i[d+1]
+                        if is_intersected(p1, p2, p3, p4): inter_count+=1
+                      if inter_count==0:
+                        head, tail = pos_i[poly_i.index(p1)], pos_i[poly_i.index(p2)]
+                        inside_edge_index[0].append(head), inside_edge_index[1].append(tail)
+
+  for i in range(len(vg_point)):
+    p1 = vg_point[i]
+    p1_id = np.count_nonzero(flag[0:i] == 2) + np.count_nonzero(flag[0:i] == 3)
+    for j in range(len(vg_point)):
+      p2 = vg_point[j]
+      if p1 == p2: continue
+      p2_id = np.count_nonzero(flag[0:j] == 2) + np.count_nonzero(flag[0:j] == 3)
+      inter_count = 0
+      for m in range(len(flag-1)):
+        if flag[m]!=1 and flag[m]!=2 and flag[m]!=3: p3, p4 = vg_point[m], vg_point[m+1]
+        if is_intersected(p1, p2, p3, p4): inter_count+=1
+      if inter_count==0:
+        head, tail = vg_point.index(p1), vg_point.index(p2)
+        cc = np.count_nonzero(flag[min(head, tail):max(head, tail)] == 2) + np.count_nonzero(flag[min(head, tail): max(head, tail)] == 3)
+        if p1_id!=p2_id and cc!=0: apart_edge_index[0].append(head), apart_edge_index[1].append(tail)
+    #print(i)
+
+  ninside_edge_index = [[],[]]
+  napart_edge_index = [[],[]]
+  exteriors = [[],[]]
+
+  if len(hole_p)!=0:
+    for i in range(len(flag)):
+      link_i = [pos_join[inside_edge_index[1][j]] for j in range(len(inside_edge_index[1])) if inside_edge_index[0][j]==i]
+      if len(link_i)==0: continue
+      ninside_edge_index[0].append(i)
+      dis_matrix = distance_matrix([pos_join[i]], link_i)
+      node_i = (link_i[np.argmin(dis_matrix[0])][0], link_i[np.argmin(dis_matrix[0])][1])
+      ninside_edge_index[1].append(vg_point.index(node_i))
+
+  for i in range(len(vg_point)-1):
+    if flag[i]!=1 and flag[i]!=2 and flag[i]!=3 : exteriors[0].append(i), exteriors[1].append(i+1)
+
+  for i in range(len(flag)):
+    link_i = [pos_join[apart_edge_index[1][j]] for j in range(len(apart_edge_index[1])) if apart_edge_index[0][j]==i]
+    if len(link_i)==0: continue
+    napart_edge_index[0].append(i)
+    dis_matrix = distance_matrix([pos_join[i]], link_i)
+    node_i = (link_i[np.argmin(dis_matrix[0])][0], link_i[np.argmin(dis_matrix[0])][1])
+    napart_edge_index[1].append(vg_point.index(node_i))
+
+  inside_edge_index, apart_edge_index = ninside_edge_index, napart_edge_index
+
+  return inside_edge_index, apart_edge_index, exteriors
+
+
+def HeteroEdge(pos,k):
+  pos_join = np.delete(pos, np.where(np.sum(pos, 1)==0)[0], axis=0)
+  pos_join = np.delete(pos_join, np.where(np.sum(pos_join, 1)==2)[0], axis=0)
+  pos_join = np.delete(pos_join, np.where(np.sum(pos_join, 1)==4)[0], axis=0)
+  flag = np.zeros(len(pos_join))
+  pos = np.delete(pos, 0, axis=0)
+  count, id = 0, 0
+  while count<k:
+    for i in range(len(pos)):
+      if pos[i][0]==0:
+        flag[i-1]=1
+        pos = np.delete(pos, i, axis=0)
+        break
+      elif pos[i][0]==1:
+        flag[i-1]=2
+        pos = np.delete(pos, i, axis=0)
+        break
+      elif pos[i][0]==2:
+        flag[i-1]=3
+        pos = np.delete(pos, i, axis=0)
+        pos_join[id:i, 0]+=count
+        count+=1
+        id = i
+        break
+  pos_join = pos_join
+  inside_edge_index, apart_edge_index, exteriors = Visi_Edge(pos_join, flag)
+
+  return pos_join, inside_edge_index, apart_edge_index, exteriors
+
+#build heterovg of k-digit from MNIST
+def NNIST_HeteroVG(df, label_df, k):
+  pos = [[0,0]]
+  label = ''
+  for i in np.random.randint(0, len(df), k):
+    while True:
+      if len(pos) == 1 and label_df[0][i] == 0: i = random.randint(0, len(df))
+      else: break
+    pos = np.append(pos, read_singe(df, i), 0)
+    pos = np.append(pos, [[2,2]], 0)
+    label = label+'%d'%(label_df[0][i])
+
+  label = int(label)
+  pos_join, inside, apart, exteriors = HeteroEdge(pos,k)
+
+  data = HeteroData()
+
+  data['vertices'].x = torch.zeros((len(pos_join), 1), dtype=torch.float)
+  data.y = torch.tensor(label, dtype=torch.int)
+  data.pos = torch.tensor(pos_join, dtype=torch.float)
+
+  data['vertices', 'inside', 'vertices'].edge_index = torch.tensor([inside[0]+inside[1]+exteriors[0],inside[1]+inside[0]+exteriors[1]], dtype=torch.long)
+  data['vertices', 'apart', 'vertices'].edge_index = torch.tensor([apart[0]+apart[1],apart[1]+apart[0]], dtype=torch.long)
+  data['vertices', 'inside', 'vertices'].edge_attr = torch.zeros((len(data['vertices', 'inside', 'vertices'].edge_index[0]),1), dtype=torch.float)
+  data['vertices', 'apart', 'vertices'].edge_attr = torch.zeros((len(data['vertices', 'apart', 'vertices'].edge_index[0]),1), dtype=torch.float)
+
+  return data
+
+mnist_filename = '/content/drive/MyDrive/MINST_Polygons/polyMNIST/mnist_polygon_test.json'
+label_filename = '/content/drive/MyDrive/MINST_Polygons/polyMNIST/mnist_label_test.json'
+df = pd.read_json(mnist_filename)
+label_df = pd.read_json(label_filename)
+
+K = 2 # number of digits
+N = 10 # number of generated graphs
+multi_mnist_dataset = []
+for k in range(2, K+1):
+  for i in tqdm(range(N)):
+    data = NNIST_HeteroVG(df, label_df, k=k)
+    multi_mnist_dataset.append(data)
+
+if not os.path.exists('/content/drive/MyDrive/MINST_Polygons/multi_mnist'):
+    os.makedirs('/content/drive/MyDrive/MINST_Polygons/multi_mnist')
+with open('/content/drive/MyDrive/MINST_Polygons/multi_mnist/multi_mnist.pkl','wb') as file:
+    pkl.dump(multi_mnist_dataset, file)

README.md

@@ -0,0 +1,3 @@
+# KDD24 PolygonGNN: Representation Learning for Polygonal Geometries with Heterogeneous Visibility Graph
+data is on [dropbox](https://www.dropbox.com/scl/fo/f7dir04pldz36n6m47m30/ABxnZk8Qyf16k0Yo75WqXpY?rlkey=f3lhgyv7um323ngpa2bmueimq&st=e4wg0uec&dl=0)
+

condainstall.txt

@@ -0,0 +1,6 @@
+conda create -n graph python=3.8
+conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch-lts -c nvidia
+conda install pytorch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 cudatoolkit=11.3 -c pytorch
+conda install pyg -c pyg
+conda install -c conda-forge pytorch_sparse
+conda install matplotlib numpy ipykernel pandas tensorboard

dataset.py

@@ -0,0 +1,213 @@
+import torch.utils.data as data
+import os
+import os.path
+import torch
+import numpy as np
+import pandas as pd
+import sys
+import pickle
+import time
+import torchvision.datasets as datasets
+import torchvision.transforms as transforms
+from PIL import Image
+from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Union
+from torchvision.datasets import VisionDataset
+from torch.utils.data import Dataset
+from datetime import date, timedelta,datetime
+import random
+import pickle as pkl
+import string
+
+valid_chars = 'EFHILOTUYZ'
+
+alphabetic_labels = [char1 + char2 for char1 in valid_chars for char2 in valid_chars]
+alphabetic_labels.sort()
+label_mapping = {label: idx for idx, label in enumerate(alphabetic_labels)} # to number
+reverse_label_mapping = {v: k for k, v in label_mapping.items()} # to alphabetic
+
+single_alphabetic_labels=[char1 for char1 in valid_chars]
+single_alphabetic_labels.sort()
+single_label_mapping = {label: idx for idx, label in enumerate(single_alphabetic_labels)}
+single_reverse_label_mapping = {v: k for k, v in single_label_mapping.items()}
+
+def get_mnist_dataset(data_dir='data/multi_mnist.pkl',Seed=0,test_ratio=0.2):
+
+    random.seed(Seed)
+    torch.manual_seed(Seed)  
+    np.random.seed(Seed) 
+
+    with open(data_dir, 'rb') as f:
+        dataset = pkl.load(f)
+    for entry in dataset:
+        entry.y -= 10
+                
+    np.random.shuffle(dataset)
+    val_test_split = int(np.around( test_ratio * len(dataset) ))
+    train_val_split = int(len(dataset)-2*val_test_split)
+    train_ds = dataset[:train_val_split]
+    val_ds = dataset[train_val_split:train_val_split+val_test_split]
+    test_ds = dataset[train_val_split+val_test_split:]  
+
+    print(data_dir)
+    print('Train: ' +str(len(train_ds)))
+    print('Val  : ' +str(len(val_ds)))
+    print('Test : ' +str(len(test_ds)))  
+        
+    return train_ds,val_ds,test_ds
+
+def get_building_dataset(data_dir='data/building_with_index.pkl',Seed=0,test_ratio=0.2):
+
+    random.seed(Seed)
+    torch.manual_seed(Seed)  
+    np.random.seed(Seed) 
+
+    with open(data_dir, 'rb') as f:
+        dataset = pkl.load(f)
+    for entry in dataset:
+        entry.y = label_mapping[entry.y]  
+         
+    np.random.shuffle(dataset)
+    val_test_split = int(np.around( test_ratio * len(dataset) ))
+    train_val_split = int(len(dataset)-2*val_test_split)
+    train_ds = dataset[:train_val_split]
+    val_ds = dataset[train_val_split:train_val_split+val_test_split]
+    test_ds = dataset[train_val_split+val_test_split:]  
+
+    print(data_dir)
+    print('Train: ' +str(len(train_ds)))
+    print('Val  : ' +str(len(val_ds)))
+    print('Test : ' +str(len(test_ds)))  
+        
+    return train_ds,val_ds,test_ds
+
+def get_mbuilding_dataset(data_dir='data/mp_building.pkl',Seed=0,test_ratio=0.2):
+
+    random.seed(Seed)
+    torch.manual_seed(Seed)  
+    np.random.seed(Seed) 
+
+    with open(data_dir, 'rb') as f:
+        dataset = pkl.load(f)
+    for entry in dataset:
+        entry.y = label_mapping[entry.y]  
+         
+    np.random.shuffle(dataset)
+    val_test_split = int(np.around( test_ratio * len(dataset) ))
+    train_val_split = int(len(dataset)-2*val_test_split)
+    train_ds = dataset[:train_val_split]
+    val_ds = dataset[train_val_split:train_val_split+val_test_split]
+    test_ds = dataset[train_val_split+val_test_split:]  
+
+    print(data_dir)
+    print('Train: ' +str(len(train_ds)))
+    print('Val  : ' +str(len(val_ds)))
+    print('Test : ' +str(len(test_ds)))  
+        
+    return train_ds,val_ds,test_ds
+
+def get_sbuilding_dataset(data_dir='data/single_building.pkl',Seed=0,test_ratio=0.2):
+
+    random.seed(Seed)
+    torch.manual_seed(Seed)  
+    np.random.seed(Seed) 
+
+    with open(data_dir, 'rb') as f:
+        dataset = pkl.load(f)
+    for entry in dataset:
+        entry.y = single_label_mapping[entry.y]  
+         
+    np.random.shuffle(dataset)
+    val_test_split = int(np.around( test_ratio * len(dataset) ))
+    train_val_split = int(len(dataset)-2*val_test_split)
+    train_ds = dataset[:train_val_split]
+    val_ds = dataset[train_val_split:train_val_split+val_test_split]
+    test_ds = dataset[train_val_split+val_test_split:]  
+
+    print(data_dir)
+    print('Train: ' +str(len(train_ds)))
+    print('Val  : ' +str(len(val_ds)))
+    print('Test : ' +str(len(test_ds)))  
+        
+    return train_ds,val_ds,test_ds
+
+def get_smnist_dataset(data_dir='data/single_mnist.pkl',Seed=0,test_ratio=0.2):
+
+    random.seed(Seed)
+    torch.manual_seed(Seed)  
+    np.random.seed(Seed) 
+
+    with open(data_dir, 'rb') as f:
+        dataset = pkl.load(f)
+         
+    np.random.shuffle(dataset)
+    val_test_split = int(np.around( test_ratio * len(dataset) ))
+    train_val_split = int(len(dataset)-2*val_test_split)
+    train_ds = dataset[:train_val_split]
+    val_ds = dataset[train_val_split:train_val_split+val_test_split]
+    test_ds = dataset[train_val_split+val_test_split:]  
+
+    print(data_dir)
+    print('Train: ' +str(len(train_ds)))
+    print('Val  : ' +str(len(val_ds)))
+    print('Test : ' +str(len(test_ds)))  
+        
+    return train_ds,val_ds,test_ds
+
+def get_dbp_dataset(data_dir='data/triple_building.pkl',Seed=0,test_ratio=0.2):
+
+    random.seed(Seed)
+    torch.manual_seed(Seed)  
+    np.random.seed(Seed) 
+
+    with open(data_dir, 'rb') as f:
+        dataset = pkl.load(f)
+    for entry in dataset:
+        entry.y = 1 if entry.y>=1 else 0
+         
+    np.random.shuffle(dataset)
+    val_test_split = int(np.around( test_ratio * len(dataset) ))
+    train_val_split = int(len(dataset)-2*val_test_split)
+    train_ds = dataset[:train_val_split]
+    val_ds = dataset[train_val_split:train_val_split+val_test_split]
+    test_ds = dataset[train_val_split+val_test_split:]  
+
+    print(data_dir)
+    print('Train: ' +str(len(train_ds)))
+    print('Val  : ' +str(len(val_ds)))
+    print('Test : ' +str(len(test_ds)))  
+        
+    return train_ds,val_ds,test_ds
+
+def affine_transform_to_range(ds, target_range=(-1, 1)):
+    # Find the extent (min and max) of coordinates in both x and y directions
+    for item in ds:
+        min_x  = torch.min(item.pos[:,0])
+        min_y  = torch.min(item.pos[:,1])
+        
+        max_x  = torch.max(item.pos[:,0])
+        max_y  = torch.max(item.pos[:,1])
+        
+        scale_x = (target_range[1] - target_range[0]) / (max_x - min_x)
+        scale_y = (target_range[1] - target_range[0]) / (max_y - min_y)
+        translate_x = target_range[0] - min_x * scale_x
+        translate_y = target_range[0] - min_y * scale_y
+
+        # Apply the affine transformation to 
+        item.pos[:,0] = item.pos[:,0] * scale_x + translate_x
+        item.pos[:,1] = item.pos[:,1] * scale_y + translate_y
+    return ds
+
+class CustomDataset(Dataset):
+    def __init__(self, data_list):
+        super(CustomDataset, self).__init__()
+        self.data_list = data_list
+    
+    def len(self):
+        return len(self.data_list)
+    
+    def get(self, idx):
+        return self.data_list[idx]
+    
+if __name__ == '__main__':
+    a,b,c=get_mnist_dataset()
+    print("")

eval.py

@@ -0,0 +1,261 @@
+import argparse
+import os
+import random
+import torch
+import pandas as pd
+import numpy as np
+import time
+import torch.optim as optim
+import scipy
+
+from matplotlib import cm
+import matplotlib.pyplot as plt
+import json
+import torch.nn.functional as F
+from torch.nn.functional import softmax
+
+torch.autograd.set_detect_anomaly(True)
+import pickle
+from torch.utils.tensorboard import SummaryWriter
+import dataset,util
+from model_new import Smodel
+import model_new
+
+
+import torch.nn as nn
+import torchvision.transforms as transforms
+import torchvision.datasets
+import torchvision.models
+import math
+import shutil
+import time
+from datetime import date, timedelta,datetime
+import torch_geometric
+from torch_geometric.data import Data, DataLoader
+from torch_geometric.nn import MessagePassing
+from torch_geometric.utils import add_self_loops
+from torch_geometric.nn import GIN,GATConv,MLP
+from torch_geometric.nn.pool import global_mean_pool,global_add_pool
+import csv
+
+blue = lambda x: '\033[94m' + x + '\033[0m'
+red = lambda x: '\033[31m' + x + '\033[0m'
+green = lambda x: '\033[32m' + x + '\033[0m'
+yellow = lambda x: '\033[33m' + x + '\033[0m'
+greenline = lambda x: '\033[42m' + x + '\033[0m'
+yellowline = lambda x: '\033[43m' + x + '\033[0m'
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model',default="our", type=str)
+    parser.add_argument('--train_batch', default=64, type=int)
+    parser.add_argument('--test_batch', default=128, type=int)
+    parser.add_argument('--share', type=str, default="0")
+    parser.add_argument('--edge_rep', type=str, default="True")
+    parser.add_argument('--batchnorm', type=str, default="True")
+    parser.add_argument('--extent_norm', type=str, default="T")
+    parser.add_argument('--spanning_tree', type=str, default="F")
+    
+    parser.add_argument('--loss_coef', default=0.1, type=float)
+    parser.add_argument('--h_ch', default=512, type=int)
+    parser.add_argument('--localdepth', type=int, default=1)
+    parser.add_argument('--num_interactions', type=int, default=4)
+    parser.add_argument('--finaldepth', type=int, default=4)
+    parser.add_argument('--classifier_depth', type=int, default=4)
+    parser.add_argument('--dropout', type=float, default=0.0)
+
+    parser.add_argument('--dataset', type=str, default='mnist')
+    parser.add_argument('--log', type=str, default="True")  
+    parser.add_argument('--test_per_round', type=int, default=10)
+    parser.add_argument('--patience', type=int, default=30)  #scheduler
+    parser.add_argument('--nepoch', type=int, default=201)
+    parser.add_argument('--lr', type=float, default=1e-4)
+    parser.add_argument('--manualSeed', type=str, default="False")
+    parser.add_argument('--man_seed', type=int, default=12345)
+    
+    parser.add_argument("--targetfiles", nargs='+', type=str, default=["Dec11-14:44:32.pth","Nov13-14:30:48.pth"])
+    args = parser.parse_args()
+    args.log=True if args.log=="True" else False
+    args.edge_rep=True if args.edge_rep=="True" else False
+    args.batchnorm=True if args.batchnorm=="True" else False
+    args.save_dir=os.path.join('./save/',args.dataset)
+    args.manualSeed=True if args.manualSeed=="True" else False
+    return args
+
+args = get_args()
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+criterion=nn.CrossEntropyLoss()
+
+def forward_HGT(args,data,model,mlpmodel):
+    data = data.to(device)
+    x,batch=data.pos, data['vertices'].batch
+    data["vertices"]['x']=data.pos
+    label=data.y.long().view(-1)
+    
+    output=model(data.x_dict, data.edge_index_dict)  
+    if args.dataset in ["dbp"]:
+        graph_embeddings=global_add_pool(output,batch)
+    else:
+        graph_embeddings=global_add_pool(output,batch)
+    graph_embeddings.clamp_(max=1e6)
+
+    output=mlpmodel(graph_embeddings)
+    # log_probs = F.log_softmax(output, dim=1)
+
+    loss = criterion(output, label) 
+    return loss,output,label, graph_embeddings
+
+def forward(args,data,model,mlpmodel):
+    data = data.to(device)
+    edge_index1=data['vertices', 'inside', 'vertices']['edge_index']
+    edge_index2=data['vertices', 'apart', 'vertices']['edge_index']
+    combined_edge_index=torch.cat([data['vertices', 'inside', 'vertices']['edge_index'],data['vertices', 'apart', 'vertices']['edge_index']],1)
+    
+    if args.spanning_tree == 'True':
+        edge_weight=torch.rand(combined_edge_index.shape[1]) + 1
+        combined_edge_index = util.build_spanning_tree_edge(combined_edge_index, edge_weight,num_nodes=num_nodes,)
+    
+    num_edge_inside=edge_index1.shape[1]
+    x,batch=data.pos, data['vertices'].batch
+    label=data.y.long().view(-1)
+    """ 
+    triplets are not the same for graphs when training 
+    """
+    num_nodes=x.shape[0]
+    edge_index_2rd, num_triplets_real, edx_jk, edx_ij = util.triplets(combined_edge_index, num_nodes)
+    
+    input_feature=torch.zeros([x.shape[0],args.h_ch],device=device) 
+    output=model(input_feature,x,[edge_index1,edge_index2], edge_index_2rd,edx_jk, edx_ij,batch,num_edge_inside,args.edge_rep)   
+    output=torch.cat(output,dim=1)
+    graph_embeddings=global_add_pool(output,batch)
+    graph_embeddings.clamp_(max=1e6)
+ 
+    output=mlpmodel(graph_embeddings)
+    # log_probs = F.log_softmax(output, dim=1)
+
+    loss = criterion(output, label) 
+    return loss,output,label,graph_embeddings
+def test(args,loader,model,mlpmodel,writer,reverse_mapping ):
+    y_hat, y_true,y_hat_logit = [], [], [],
+    embeddings=[]
+
+    loss_total, pred_num = 0, 0
+    model.eval()
+    mlpmodel.eval()
+    with torch.no_grad():
+        for data in loader:
+            if args.model=="our":
+                loss,output,label,embedding  =forward(args,data,model,mlpmodel)
+            elif args.model in ["HGT","HAN"]:
+                loss,output,label,embedding =forward_HGT(args,data,model,mlpmodel)            
+            _, pred = output.topk(1, dim=1, largest=True, sorted=True)
+            pred,label,output=pred.cpu(),label.cpu(),output.cpu()
+            y_hat += list(pred.detach().numpy().reshape(-1))
+            y_true += list(label.detach().numpy().reshape(-1))
+            y_hat_logit+=list(output.detach().numpy())
+            embeddings.append(embedding)
+            
+            pred_num += len(label.reshape(-1, 1))
+            loss_total += loss.detach() * len(label.reshape(-1, 1))
+            
+    y_true_str=[reverse_mapping(item) for item in y_true] 
+    writer.add_embedding(torch.cat(embeddings,dim=0).detach().cpu(),metadata=y_true_str,tag="numbers")
+    writer.close()      
+    return loss_total/pred_num,y_hat, y_true, y_hat_logit  
+ 
+def main(args,train_Loader,val_Loader,test_Loader):
+    donefiles=os.listdir(os.path.join(args.save_dir,args.model,'model'))
+    tensorboard_dir=os.path.join(args.save_dir,args.model,'log') 
+    if args.dataset in ["mnist","mnist_sparse"]:
+        reverse_mapping=lambda x: x + 10
+        # list(map(lambda x: x - 10, []))
+    elif args.dataset in ["building","mbuilding"]:
+        reverse_mapping=lambda x: dataset.reverse_label_mapping[x]
+    elif args.dataset in ["sbuilding"]:
+        reverse_mapping=lambda x: dataset.single_reverse_label_mapping[x]
+    elif args.dataset in ["dbp","smnist"]:
+        reverse_mapping=lambda x: x
+    for file in donefiles:
+        if file not in args.targetfiles:
+            continue
+        else:
+            print(file)
+            saved_dict=torch.load(os.path.join(args.save_dir,args.model,'model',file))   
+            if saved_dict['args'].dataset in ["mnist","mnist_sparse"]:
+                x_out=90
+            elif saved_dict['args'].dataset in ["building","mbuilding"]:
+                x_out=100
+            elif saved_dict['args'].dataset in ["sbuilding","smnist"]:
+                x_out=10
+            elif saved_dict['args'].dataset in ["dbp"]:
+                x_out=2
+            if saved_dict['args'].model=="our":
+                model=Smodel(h_channel=saved_dict['args'].h_ch,input_featuresize=saved_dict['args'].h_ch,\
+                                localdepth=saved_dict['args'].localdepth,num_interactions=saved_dict['args'].num_interactions,finaldepth=saved_dict['args'].finaldepth,share=saved_dict['args'].share,batchnorm=saved_dict['args'].batchnorm)
+                mlpmodel=MLP(in_channels=saved_dict['args'].h_ch*saved_dict['args'].num_interactions, hidden_channels=saved_dict['args'].h_ch,out_channels=x_out, num_layers=saved_dict['args'].classifier_depth)
+            elif saved_dict['args'].model=="HGT":
+                model=model_new.HGT(hidden_channels=saved_dict['args'].h_ch, out_channels=saved_dict['args'].h_ch, num_heads=2, num_layers=saved_dict['args'].num_interactions)
+                mlpmodel=MLP(in_channels=saved_dict['args'].h_ch, hidden_channels=saved_dict['args'].h_ch,out_channels=x_out, num_layers=saved_dict['args'].classifier_depth,dropout=saved_dict['args'].dropout)
+            elif saved_dict['args'].model=="HAN":
+                model=model_new.HAN(hidden_channels=saved_dict['args'].h_ch, out_channels=saved_dict['args'].h_ch, num_heads=2, num_layers=saved_dict['args'].num_interactions)
+                mlpmodel=MLP(in_channels=saved_dict['args'].h_ch, hidden_channels=saved_dict['args'].h_ch,out_channels=x_out, num_layers=saved_dict['args'].classifier_depth,dropout=saved_dict['args'].dropout)
+            model.to(device), mlpmodel.to(device)         
+            try:
+                model.load_state_dict(saved_dict['model'],strict=True)
+                mlpmodel.load_state_dict(saved_dict['mlpmodel'],strict=True)
+            except OSError:
+                print('loadfail: ',file)
+                pass
+            print(saved_dict['args'])
+
+            writer = SummaryWriter(os.path.join(tensorboard_dir,file+"_embedding"))
+            test_loss, yhat_test, ytrue_test, yhatlogit_test = test(saved_dict['args'],test_Loader,model,mlpmodel,writer,reverse_mapping)
+                
+            pred_dir=os.path.join(tensorboard_dir,file+"_test_record")
+            to_save_dict={'labels':ytrue_test,'yhat':yhat_test,'yhat_logit':yhatlogit_test}
+            torch.save(to_save_dict, pred_dir)
+                           
+            test_acc=util.calculate(yhat_test,ytrue_test,yhatlogit_test)
+            util.print_1(0,'Test', {"loss":test_loss,"acc":test_acc},color=blue) 
+
+        
+if __name__ == '__main__':
+    Seed = 0
+    test_ratio=0.2
+    print("data splitting Random Seed: ", Seed)
+    if args.dataset in ["mnist"]:
+        args.data_dir='data/multi_mnist_with_index.pkl'
+        train_ds,val_ds,test_ds=dataset.get_mnist_dataset(args.data_dir,Seed,test_ratio=test_ratio)
+    elif args.dataset in ["mnist_sparse"]:
+        args.data_dir='data/multi_mnist_sparse.pkl'  
+        train_ds,val_ds,test_ds=dataset.get_mnist_dataset(args.data_dir,Seed,test_ratio=test_ratio)      
+    elif args.dataset in ["building"]:
+        args.data_dir='data/building_with_index.pkl'
+        train_ds,val_ds,test_ds=dataset.get_mbuilding_dataset(args.data_dir,Seed,test_ratio=test_ratio)
+    elif args.dataset in ["mbuilding"]:
+        args.data_dir='data/mp_building.pkl'
+        train_ds,val_ds,test_ds=dataset.get_building_dataset(args.data_dir,Seed,test_ratio=test_ratio)
+    elif args.dataset in ["sbuilding"]:
+        args.data_dir='data/single_building.pkl'
+        train_ds,val_ds,test_ds=dataset.get_sbuilding_dataset(args.data_dir,Seed,test_ratio=test_ratio)
+    elif args.dataset in ["smnist"]:
+        args.data_dir='data/single_mnist.pkl'
+        train_ds,val_ds,test_ds=dataset.get_smnist_dataset(args.data_dir,Seed,test_ratio=test_ratio)
+    elif args.dataset in ['dbp']:
+        args.data_dir='data/triple_building_600.pkl'
+        train_ds,val_ds,test_ds=dataset.get_dbp_dataset(args.data_dir,Seed,test_ratio=test_ratio)
+
+    if args.extent_norm=="T":
+        train_ds= dataset.affine_transform_to_range(train_ds,target_range=(-1, 1))
+        val_ds= dataset.affine_transform_to_range(val_ds,target_range=(-1, 1))
+        test_ds= dataset.affine_transform_to_range(test_ds,target_range=(-1, 1))              
+    train_loader = torch_geometric.loader.DataLoader(train_ds,batch_size=args.train_batch, shuffle=False,pin_memory=True) 
+    val_loader = torch_geometric.loader.DataLoader(val_ds, batch_size=args.test_batch, shuffle=False, pin_memory=True)
+    test_loader = torch_geometric.loader.DataLoader(test_ds,batch_size=args.test_batch, shuffle=False,pin_memory=True)
+
+    Seed=random.randint(1, 10000)
+    print("Random Seed: ", Seed)
+    random.seed(Seed)
+    torch.manual_seed(Seed)  
+    np.random.seed(Seed)     
+    main(args,train_loader,val_loader,test_loader)

model_new.py

@@ -0,0 +1,224 @@
+import math
+from math import pi as PI
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.parallel
+import torch.utils.data
+import torch_geometric.transforms as T
+from torch.nn import ModuleList, Parameter
+from torch_geometric.nn import HANConv, HEATConv, HGTConv, Linear
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.nn.dense.linear import Linear
+# from dataset import 
+from torch_geometric.nn.inits import glorot, zeros
+from torch_geometric.utils import softmax
+from torch_scatter import scatter
+
+from util import get_angle, get_theta, triplets
+
+class Smodel(nn.Module):
+    def __init__(self,  h_channel=16,input_featuresize=32,localdepth=2,num_interactions=3,finaldepth=3,share='0',batchnorm="True"):
+        super(Smodel,self).__init__()
+        self.training=True
+        self.h_channel = h_channel
+        self.input_featuresize=input_featuresize
+        self.localdepth = localdepth
+        self.num_interactions=num_interactions
+        self.finaldepth=finaldepth
+        self.batchnorm = batchnorm        
+        self.activation=nn.ReLU()
+        self.att = Parameter(torch.ones(4),requires_grad=True)
+
+        num_gaussians=(1,1,1)
+        self.mlp_geo = ModuleList()
+        for i in range(self.localdepth):
+            if i == 0:
+                self.mlp_geo.append(Linear(sum(num_gaussians), h_channel))
+            else:
+                self.mlp_geo.append(Linear(h_channel, h_channel))
+            if self.batchnorm == "True":
+                self.mlp_geo.append(nn.BatchNorm1d(h_channel))
+            self.mlp_geo.append(self.activation)            
+         
+        self.mlp_geo_backup = ModuleList()
+        for i in range(self.localdepth):
+            if i == 0:
+                self.mlp_geo_backup.append(Linear(4, h_channel))
+            else:
+                self.mlp_geo_backup.append(Linear(h_channel, h_channel))
+            if self.batchnorm == "True":
+                self.mlp_geo_backup.append(nn.BatchNorm1d(h_channel))
+            self.mlp_geo_backup.append(self.activation)        
+        self.translinear=Linear(input_featuresize+1, self.h_channel)
+        self.interactions= ModuleList()
+        for i in range(self.num_interactions):
+            block = SPNN(
+                in_ch=self.input_featuresize,
+                hidden_channels=self.h_channel,
+                activation=self.activation,
+                finaldepth=self.finaldepth,
+                batchnorm=self.batchnorm,
+                num_input_geofeature=self.h_channel
+            )
+            self.interactions.append(block)
+        self.reset_parameters()
+    def reset_parameters(self):
+        for lin in self.mlp_geo:
+            if isinstance(lin, Linear):
+                torch.nn.init.xavier_uniform_(lin.weight)
+                lin.bias.data.fill_(0)
+        for i in (self.interactions):
+                i.reset_parameters()
+
+    def single_forward(self, input_feature,coords,edge_index,edge_index_2rd, edx_jk, edx_ij,batch,num_edge_inside,edge_rep):
+        if edge_rep:
+            i, j, k = edge_index_2rd
+            edge_index1,edge_index2=  edge_index
+            edge_index_all=torch.cat([edge_index1,edge_index2],1)
+            distance_ij=(coords[j] - coords[i]).norm(p=2, dim=1)
+            distance_jk=(coords[j] - coords[k]).norm(p=2, dim=1)
+            theta_ijk = get_angle(coords[j] - coords[i], coords[k] - coords[j])
+            geo_encoding_1st=distance_ij[:,None]
+            geo_encoding=torch.cat([geo_encoding_1st,distance_jk[:,None],theta_ijk[:,None]],dim=-1)
+        else:    
+            coords_j = coords[edge_index[0]]
+            coords_i = coords[edge_index[1]]
+            geo_encoding=torch.cat([coords_j,coords_i],dim=-1)
+        if edge_rep:
+            for lin in self.mlp_geo:
+                geo_encoding=lin(geo_encoding)
+        else:
+            for lin in self.mlp_geo_backup:
+                geo_encoding=lin(geo_encoding)
+            geo_encoding=torch.zeros_like(geo_encoding,device=geo_encoding.device,dtype=geo_encoding.dtype)
+        node_feature= input_feature
+        node_feature_list=[]
+        for interaction in self.interactions:
+            node_feature =  interaction(node_feature,geo_encoding,edge_index_2rd,edx_jk,edx_ij,num_edge_inside,self.att)
+            node_feature_list.append(node_feature)
+        return node_feature_list
+    def forward(self, input_feature, coords,edge_index,edge_index_2rd, edx_jk, edx_ij,batch,num_edge_inside,edge_rep):
+        output=self.single_forward(input_feature,coords,edge_index,edge_index_2rd, edx_jk, edx_ij,batch,num_edge_inside,edge_rep)
+        return output
+    
+class SPNN(torch.nn.Module):
+    def __init__(
+        self,
+        in_ch,
+        hidden_channels,
+        activation=torch.nn.ReLU(),
+        finaldepth=3,
+        batchnorm="True",
+        num_input_geofeature=13
+    ):
+        super(SPNN, self).__init__()
+        self.activation = activation
+        self.finaldepth = finaldepth
+        self.batchnorm = batchnorm
+        self.num_input_geofeature=num_input_geofeature
+        
+        self.WMLP_list = ModuleList()
+        for _ in range(4):
+            WMLP = ModuleList()
+            for i in range(self.finaldepth + 1):
+                if i == 0:
+                    WMLP.append(Linear(hidden_channels*3+num_input_geofeature, hidden_channels))
+                else:
+                    WMLP.append(Linear(hidden_channels, hidden_channels))  
+                if self.batchnorm == "True":
+                    WMLP.append(nn.BatchNorm1d(hidden_channels))
+                WMLP.append(self.activation)
+            self.WMLP_list.append(WMLP)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for mlp in self.WMLP_list:
+            for lin in mlp:
+                if isinstance(lin, Linear):
+                    torch.nn.init.xavier_uniform_(lin.weight)
+                    lin.bias.data.fill_(0)
+    def forward(self, node_feature,geo_encoding,edge_index_2rd,edx_jk,edx_ij,num_edge_inside,att):
+        i,j,k = edge_index_2rd
+        if node_feature is None:
+            concatenated_vector = geo_encoding
+        else:
+            node_attr_0st = node_feature[i]
+            node_attr_1st = node_feature[j]
+            node_attr_2 = node_feature[k]
+            concatenated_vector = torch.cat(
+                [
+                    node_attr_0st,
+                    node_attr_1st,node_attr_2,
+                    geo_encoding,
+                ],
+                dim=-1,
+            )
+        x_i = concatenated_vector
+        
+        edge1_edge1_mask = (edx_ij < num_edge_inside) & (edx_jk < num_edge_inside) 
+        edge1_edge2_mask = (edx_ij < num_edge_inside) & (edx_jk >= num_edge_inside)
+        edge2_edge1_mask = (edx_ij >= num_edge_inside) & (edx_jk < num_edge_inside)
+        edge2_edge2_mask = (edx_ij >= num_edge_inside) & (edx_jk >= num_edge_inside)  
+        masks=[edge1_edge1_mask,edge1_edge2_mask,edge2_edge1_mask,edge2_edge2_mask]
+        
+        x_output=torch.zeros(x_i.shape[0],self.WMLP_list[0][0].weight.shape[0],device=x_i.device)
+        for index in range(4):
+            WMLP=self.WMLP_list[index]
+            x=x_i[masks[index]]
+            for lin in WMLP:
+                x=lin(x)    
+            x = F.leaky_relu(x)*att[index]
+            x_output[masks[index]]+=x
+    
+        out_feature = scatter(x_output, i, dim=0, reduce='add')  
+        return out_feature
+
+class HGT(torch.nn.Module):
+    def __init__(self, hidden_channels, out_channels, num_heads, num_layers):
+        super().__init__()
+
+        self.lin_dict = torch.nn.ModuleDict()
+        for node_type in ["vertices"]:
+            self.lin_dict[node_type] = Linear(-1, hidden_channels)
+
+        self.convs = torch.nn.ModuleList()
+        for _ in range(num_layers):
+            conv = HGTConv(hidden_channels, hidden_channels, (['vertices'],[('vertices', 'inside', 'vertices'), ('vertices', 'apart', 'vertices')]),
+                           num_heads, group='sum')
+            self.convs.append(conv)
+
+        self.lin = Linear(hidden_channels, out_channels)
+
+    def forward(self, x_dict, edge_index_dict):
+        for node_type, x in x_dict.items():
+            x_dict[node_type]=self.lin_dict[node_type](x).relu_()
+
+        for conv in self.convs:
+            x_dict = conv(x_dict, edge_index_dict)
+        return self.lin(x_dict['vertices'])
+class HAN(torch.nn.Module):
+    def __init__(self, hidden_channels, out_channels, num_heads, num_layers):
+        super().__init__()
+
+        self.lin_dict = torch.nn.ModuleDict()
+        for node_type in ["vertices"]:
+            self.lin_dict[node_type] = Linear(-1, hidden_channels)
+
+        self.convs = torch.nn.ModuleList()
+        for _ in range(num_layers):
+            conv = HANConv(hidden_channels, hidden_channels, (['vertices'],[('vertices', 'inside', 'vertices'), ('vertices', 'apart', 'vertices')]),
+                           num_heads)
+            self.convs.append(conv)
+
+        self.lin = Linear(hidden_channels, out_channels)
+
+    def forward(self, x_dict, edge_index_dict):
+        for node_type, x in x_dict.items():
+            x_dict[node_type]=self.lin_dict[node_type](x).relu_()
+
+        for conv in self.convs:
+            x_dict = conv(x_dict, edge_index_dict)
+        return self.lin(x_dict['vertices'])

train_new.py

@@ -0,0 +1,397 @@
+import argparse
+import copy
+import json
+import os
+import random
+import time
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import scipy
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
+from matplotlib import cm
+from sklearn.metrics import (auc, explained_variance_score, f1_score,
+                             mean_absolute_error, mean_squared_error,
+                             precision_score, r2_score, recall_score,
+                             roc_auc_score, roc_curve)
+from torch.nn.functional import softmax
+from torch_geometric.utils import subgraph
+
+torch.autograd.set_detect_anomaly(True)
+import math
+import pickle
+import time
+from datetime import date, datetime, timedelta
+
+import torch.nn as nn
+import torch_geometric
+import torchvision.datasets
+import torchvision.models
+import torchvision.transforms as transforms
+from torch.utils.tensorboard import SummaryWriter
+from torch_geometric.nn import GIN, MLP, GATConv
+from torch_geometric.nn.pool import global_add_pool, global_mean_pool
+from torch_geometric.utils import add_self_loops
+
+import dataset
+import model_new
+import util
+from dataset import label_mapping, reverse_label_mapping
+from model_new import Smodel
+
+blue = lambda x: '\033[94m' + x + '\033[0m'
+red = lambda x: '\033[31m' + x + '\033[0m'
+green = lambda x: '\033[32m' + x + '\033[0m'
+yellow = lambda x: '\033[33m' + x + '\033[0m'
+greenline = lambda x: '\033[42m' + x + '\033[0m'
+yellowline = lambda x: '\033[43m' + x + '\033[0m'
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model',default="our", type=str)
+    parser.add_argument('--train_batch', default=64, type=int)
+    parser.add_argument('--test_batch', default=128, type=int)
+    parser.add_argument('--share', type=str, default="0")
+    parser.add_argument('--edge_rep', type=str, default="True")
+    parser.add_argument('--batchnorm', type=str, default="True")
+    parser.add_argument('--extent_norm', type=str, default="T")
+    parser.add_argument('--spanning_tree', type=str, default="T")
+    
+    parser.add_argument('--loss_coef', default=0.1, type=float)
+    parser.add_argument('--h_ch', default=512, type=int)
+    parser.add_argument('--localdepth', type=int, default=1)
+    parser.add_argument('--num_interactions', type=int, default=4)
+    parser.add_argument('--finaldepth', type=int, default=4)
+    parser.add_argument('--classifier_depth', type=int, default=4)
+    parser.add_argument('--dropout', type=float, default=0.0)
+
+    parser.add_argument('--dataset', type=str, default='mnist')
+    parser.add_argument('--log', type=str, default="True")   
+    parser.add_argument('--test_per_round', type=int, default=10)
+    parser.add_argument('--patience', type=int, default=30)  #scheduler
+    parser.add_argument('--nepoch', type=int, default=301)
+    parser.add_argument('--lr', type=float, default=1e-4)
+    parser.add_argument('--manualSeed', type=str, default="False")
+    parser.add_argument('--man_seed', type=int, default=12345)
+    args = parser.parse_args()
+    args.log=True if args.log=="True" else False 
+    args.edge_rep=True if args.edge_rep=="True" else False
+    args.batchnorm=True if args.batchnorm=="True" else False
+    args.save_dir=os.path.join('./save/',args.dataset,args.model)
+    args.manualSeed=True if args.manualSeed=="True" else False
+    return args
+
+args = get_args()
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+criterion=nn.CrossEntropyLoss()
+if args.dataset in ["mnist"]:
+    x_out=90
+    args.data_dir='data/multi_mnist_with_index.pkl'
+elif args.dataset in ["mnist_sparse"]:
+    x_out=90
+    args.data_dir='data/multi_mnist_sparse.pkl'
+elif args.dataset in ["building"]:
+    x_out=100
+    args.data_dir='data/building_with_index.pkl'
+elif args.dataset in ["mbuilding"]:
+    x_out=100
+    args.data_dir='data/mp_building.pkl'
+elif args.dataset in ["sbuilding"]:
+    x_out=10
+    args.data_dir='data/single_building.pkl'
+elif args.dataset in ["smnist"]:
+    x_out=10
+    args.data_dir='data/single_mnist.pkl'
+elif args.dataset in ["dbp"]:
+    x_out=2
+    args.data_dir='data/triple_building_600.pkl'
+    
+    
+if args.model=="our":
+    model=Smodel(h_channel=args.h_ch,input_featuresize=args.h_ch,\
+                    localdepth=args.localdepth,num_interactions=args.num_interactions,finaldepth=args.finaldepth,share=args.share,batchnorm=args.batchnorm)
+    mlpmodel=MLP(in_channels=args.h_ch*args.num_interactions, hidden_channels=args.h_ch,out_channels=x_out, num_layers=args.classifier_depth,dropout=args.dropout)
+
+elif args.model=="HGT":
+    model=model_new.HGT(hidden_channels=args.h_ch, out_channels=args.h_ch, num_heads=2, num_layers=args.num_interactions)
+    mlpmodel=MLP(in_channels=args.h_ch, hidden_channels=args.h_ch,out_channels=x_out, num_layers=args.classifier_depth,dropout=args.dropout)
+elif args.model=="HAN":
+    model=model_new.HAN(hidden_channels=args.h_ch, out_channels=args.h_ch, num_heads=2, num_layers=args.num_interactions)
+    mlpmodel=MLP(in_channels=args.h_ch, hidden_channels=args.h_ch,out_channels=x_out, num_layers=args.classifier_depth,dropout=args.dropout)
+
+model.to(device), mlpmodel.to(device)
+opt_list=list(model.parameters())+list(mlpmodel.parameters())
+
+optimizer = torch.optim.Adam( opt_list, lr=args.lr)    
+scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.1, patience=args.patience, min_lr=1e-8)   
+
+def contrastive_loss(embeddings,labels,margin):
+    
+    positive_mask = labels.view(-1, 1) == labels.view(1, -1)
+    negative_mask = ~positive_mask
+
+    # Calculate the number of positive and negative pairs
+    num_positive_pairs = positive_mask.sum() - labels.shape[0] 
+    num_negative_pairs = negative_mask.sum()
+
+    # If there are no negative pairs, return a placeholder loss
+    if num_negative_pairs==0 or num_positive_pairs== 0:
+        print("all pos or neg")
+        return torch.tensor(0, dtype=torch.float)
+    # Calculate the pairwise Euclidean distances between embeddings
+    distances = torch.cdist(embeddings, embeddings)/np.sqrt(embeddings.shape[1])
+    
+    if num_positive_pairs>num_negative_pairs:
+        # Sample an equal number of + pairs 
+        positive_indices = torch.nonzero(positive_mask)
+        random_positive_indices = torch.randperm(len(positive_indices))[:num_negative_pairs]
+        selected_positive_indices = positive_indices[random_positive_indices]
+
+        # Select corresponding negative pairs
+        negative_mask.fill_diagonal_(False)
+        negative_distances = distances[negative_mask].view(-1, 1)
+        positive_distances = distances[selected_positive_indices[:,0],selected_positive_indices[:,1]].view(-1, 1)
+    else: # case for most datasets
+        # Sample an equal number of - pairs 
+        negative_indices = torch.nonzero(negative_mask)
+        random_negative_indices = torch.randperm(len(negative_indices))[:num_positive_pairs]
+        selected_negative_indices = negative_indices[random_negative_indices]
+
+        # Select corresponding positive pairs
+        positive_mask.fill_diagonal_(False)
+        positive_distances = distances[positive_mask].view(-1, 1)
+        negative_distances = distances[selected_negative_indices[:,0],selected_negative_indices[:,1]].view(-1, 1)
+
+    # Calculate the loss for positive and negative pairs
+    loss = (positive_distances - negative_distances + margin).clamp(min=0).mean()
+    return loss
+
+def forward_HGT(data,model,mlpmodel):
+    data = data.to(device)
+    x,batch=data.pos, data['vertices'].batch
+    data["vertices"]['x']=data.pos
+    label=data.y.long().view(-1)
+
+    optimizer.zero_grad()
+    
+    output=model(data.x_dict, data.edge_index_dict)  
+    if args.dataset in ["dbp"]:
+        graph_embeddings=global_add_pool(output,batch)
+    else:
+        graph_embeddings=global_add_pool(output,batch)
+    graph_embeddings.clamp_(max=1e6)
+    c_loss=contrastive_loss(graph_embeddings,label,margin=1)
+    output=mlpmodel(graph_embeddings)
+    # log_probs = F.log_softmax(output, dim=1)
+
+    loss = criterion(output, label) 
+    loss+=c_loss*args.loss_coef
+    return loss,c_loss*args.loss_coef,output,label    
+    
+def forward(data,model,mlpmodel):
+    data = data.to(device)
+    edge_index1=data['vertices', 'inside', 'vertices']['edge_index']
+    edge_index2=data['vertices', 'apart', 'vertices']['edge_index']
+    combined_edge_index=torch.cat([data['vertices', 'inside', 'vertices']['edge_index'],data['vertices', 'apart', 'vertices']['edge_index']],1)
+    num_edge_inside=edge_index1.shape[1]
+    
+    if args.spanning_tree == 'T':
+        edge_weight=torch.rand(combined_edge_index.shape[1]) + 1
+        undirected_spanning_edge = util.build_spanning_tree_edge(combined_edge_index, edge_weight,num_nodes=data.pos.shape[0])
+        
+        edge_set_1 = set(map(tuple, edge_index2.t().tolist()))
+        edge_set_2 = set(map(tuple, undirected_spanning_edge.t().tolist()))
+
+        common_edges = edge_set_1.intersection(edge_set_2)
+        common_edges_tensor = torch.tensor(list(common_edges), dtype=torch.long).t().to(device)
+        spanning_edge=torch.cat([edge_index1,common_edges_tensor],1)
+        combined_edge_index=spanning_edge
+    x,batch=data.pos, data['vertices'].batch
+    label=data.y.long().view(-1)
+
+    num_nodes=x.shape[0]
+    edge_index_2rd, num_triplets_real, edx_jk, edx_ij = util.triplets(combined_edge_index, num_nodes)
+    optimizer.zero_grad()
+    input_feature=torch.zeros([x.shape[0],args.h_ch],device=device) 
+    output=model(input_feature,x,[edge_index1,edge_index2], edge_index_2rd,edx_jk, edx_ij,batch,num_edge_inside,args.edge_rep)  
+    output=torch.cat(output,dim=1)
+    if args.dataset in ["dbp"]:
+        graph_embeddings=global_add_pool(output,batch)
+    else:
+        graph_embeddings=global_add_pool(output,batch)
+    graph_embeddings.clamp_(max=1e6)
+    c_loss=contrastive_loss(graph_embeddings,label,margin=1)
+    output=mlpmodel(graph_embeddings)
+
+    loss = criterion(output, label) 
+    loss+=c_loss*args.loss_coef
+    return loss,c_loss*args.loss_coef,output,label
+def train(train_Loader,model,mlpmodel ):
+    epochloss=0
+    epochcloss=0
+    y_hat, y_true,y_hat_logit = [], [], []        
+    optimizer.zero_grad()
+    model.train()
+    mlpmodel.train()
+    for i,data in enumerate(train_Loader):
+        if args.model=="our":
+            loss,c_loss,output,label  =forward(data,model,mlpmodel)
+        elif args.model in ["HGT","HAN"]:
+            loss,c_loss,output,label  =forward_HGT(data,model,mlpmodel)
+
+        loss.backward()
+        optimizer.step()
+        epochloss+=loss.detach().cpu()
+        epochcloss+=c_loss.detach().cpu()
+        
+        _, pred = output.topk(1, dim=1, largest=True, sorted=True)
+        pred,label,output=pred.cpu(),label.cpu(),output.cpu()
+        y_hat += list(pred.detach().numpy().reshape(-1))
+        y_true += list(label.detach().numpy().reshape(-1))
+        y_hat_logit+=list(output.detach().numpy())
+    return epochloss.item()/len(train_Loader),epochcloss.item()/len(train_Loader),y_hat, y_true,y_hat_logit
+
+def test(loader,model,mlpmodel ):
+    y_hat, y_true,y_hat_logit = [], [], []
+    loss_total, pred_num = 0, 0
+    model.eval()
+    mlpmodel.eval()
+    with torch.no_grad():
+        for data in loader:
+            if args.model=="our":
+                loss,c_loss,output,label  =forward(data,model,mlpmodel)
+            elif args.model in ["HGT","HAN"]:
+                loss,c_loss,output,label  =forward_HGT(data,model,mlpmodel)
+            
+            _, pred = output.topk(1, dim=1, largest=True, sorted=True)
+            pred,label,output=pred.cpu(),label.cpu(),output.cpu()
+            y_hat += list(pred.detach().numpy().reshape(-1))
+            y_true += list(label.detach().numpy().reshape(-1))
+            y_hat_logit+=list(output.detach().numpy())
+            
+            pred_num += len(label.reshape(-1, 1))
+            loss_total += loss.detach() * len(label.reshape(-1, 1))
+    return loss_total/pred_num,y_hat, y_true, y_hat_logit
+def main(args,train_Loader,val_Loader,test_Loader):
+    best_val_trigger = -1
+    old_lr=1e3
+    suffix="{}{}-{}:{}:{}".format(datetime.now().strftime("%h"),
+                                    datetime.now().strftime("%d"),
+                                    datetime.now().strftime("%H"),
+                                    datetime.now().strftime("%M"),
+                                    datetime.now().strftime("%S"))        
+    if args.log: writer = SummaryWriter(os.path.join(tensorboard_dir,suffix))
+
+    for epoch in range(args.nepoch):
+        train_loss,train_closs,y_hat, y_true,y_hat_logit=train(train_Loader,model,mlpmodel )
+
+        train_acc=util.calculate(y_hat,y_true,y_hat_logit)
+        try:util.record({"loss":train_loss,"closs":train_closs,"acc":train_acc},epoch,writer,"Train") 
+        except: pass
+        util.print_1(epoch,'Train',{"loss":train_loss,"closs":train_closs,"acc":train_acc})
+        if epoch % args.test_per_round == 0:
+            val_loss, yhat_val, ytrue_val, yhatlogit_val = test(val_Loader,model,mlpmodel )
+            test_loss, yhat_test, ytrue_test, yhatlogit_test = test(test_Loader,model,mlpmodel )
+            val_acc=util.calculate(yhat_val,ytrue_val,yhatlogit_val)
+            try:util.record({"loss":val_loss,"acc":val_acc},epoch,writer,"Val")
+            except: pass
+            util.print_1(epoch,'Val',{"loss":val_loss,"acc":val_acc},color=blue) 
+            test_acc=util.calculate(yhat_test,ytrue_test,yhatlogit_test)
+            try:util.record({"loss":test_loss,"acc":test_acc},epoch,writer,"Test")            
+            except: pass
+            util.print_1(epoch,'Test',{"loss":test_loss,"acc":test_acc},color=blue)
+            val_trigger=val_acc
+            if val_trigger > best_val_trigger:
+                best_val_trigger = val_trigger
+                best_model = copy.deepcopy(model)
+                best_mlpmodel=copy.deepcopy(mlpmodel)
+                best_info=[epoch,val_trigger]
+        """ 
+        update lr when epoch≥30
+        """
+        if epoch >= 30:
+            lr = scheduler.optimizer.param_groups[0]['lr']
+            if old_lr!=lr:
+                print(red('lr'), epoch, (lr), sep=', ')
+                old_lr=lr
+            scheduler.step(val_trigger)        
+    """
+    use best model to get best model result 
+    """
+    val_loss, yhat_val, ytrue_val, yhat_logit_val  = test(val_Loader,best_model,best_mlpmodel)
+    test_loss, yhat_test, ytrue_test, yhat_logit_test= test(test_Loader,best_model,best_mlpmodel)
+
+    val_acc=util.calculate(yhat_val,ytrue_val,yhat_logit_val)
+    util.print_1(best_info[0],'BestVal',{"loss":val_loss,"acc":val_acc},color=blue)
+    test_acc=util.calculate(yhat_test,ytrue_test,yhat_logit_test)
+    util.print_1(best_info[0],'BestTest',{"loss":test_loss,"acc":test_acc},color=blue)
+    if args.model=="our":print(best_model.att)
+                                                            
+    """
+    save training info and best result 
+    """
+    result_file=os.path.join(info_dir, suffix)
+    with open(result_file, 'w') as f:
+        print("Random Seed: ", Seed,file=f)
+        print(f"acc  val : {val_acc:.3f}, Test : {test_acc:.3f}", file=f)
+        print(f"Best info: {best_info}", file=f)
+        for i in [[a,getattr(args, a)] for a in args.__dict__]:
+            print(i,sep='\n',file=f)
+    to_save_dict={'model':best_model.state_dict(),'mlpmodel':best_mlpmodel.state_dict(),'args':args,'labels':ytrue_test,'yhat':yhat_test,'yhat_logit':yhat_logit_test}
+    torch.save(to_save_dict, os.path.join(model_dir,suffix+'.pth') )
+    print("done")
+
+if __name__ == '__main__':
+    """
+    build dir 
+    """
+    if not os.path.exists(args.save_dir):
+        os.makedirs(args.save_dir,exist_ok=True)
+    tensorboard_dir=os.path.join(args.save_dir,'log')
+    if not os.path.exists(tensorboard_dir):
+        os.makedirs(tensorboard_dir,exist_ok=True)
+    model_dir=os.path.join(args.save_dir,'model')
+    if not os.path.exists(model_dir):
+        os.makedirs(model_dir,exist_ok=True)    
+    info_dir=os.path.join(args.save_dir,'info')
+    if not os.path.exists(info_dir):
+        os.makedirs(info_dir,exist_ok=True)      
+
+    Seed = 0
+    test_ratio=0.2
+    print("data splitting Random Seed: ", Seed)
+    if args.dataset in ['mnist',"mnist_sparse"]:
+        train_ds,val_ds,test_ds=dataset.get_mnist_dataset(args.data_dir,Seed,test_ratio=test_ratio)
+    elif args.dataset in ['building']:
+        train_ds,val_ds,test_ds=dataset.get_building_dataset(args.data_dir,Seed,test_ratio=test_ratio)
+    elif args.dataset in ['mbuilding']:
+        train_ds,val_ds,test_ds=dataset.get_mbuilding_dataset(args.data_dir,Seed,test_ratio=test_ratio)
+    elif args.dataset in ['sbuilding']:
+        train_ds,val_ds,test_ds=dataset.get_sbuilding_dataset(args.data_dir,Seed,test_ratio=test_ratio)
+    elif args.dataset in ['smnist']:
+        train_ds,val_ds,test_ds=dataset.get_smnist_dataset(args.data_dir,Seed,test_ratio=test_ratio)    
+    elif args.dataset in ['dbp']:
+        train_ds,val_ds,test_ds=dataset.get_dbp_dataset(args.data_dir,Seed,test_ratio=test_ratio)
+    if args.extent_norm=="T":
+        train_ds= dataset.affine_transform_to_range(train_ds,target_range=(-1, 1))
+        val_ds= dataset.affine_transform_to_range(val_ds,target_range=(-1, 1))
+        test_ds= dataset.affine_transform_to_range(test_ds,target_range=(-1, 1))
+                
+    train_loader = torch_geometric.loader.DataLoader(train_ds,batch_size=args.train_batch, shuffle=False,pin_memory=True,drop_last=True) 
+    val_loader = torch_geometric.loader.DataLoader(val_ds, batch_size=args.test_batch, shuffle=False, pin_memory=True)
+    test_loader = torch_geometric.loader.DataLoader(test_ds,batch_size=args.test_batch, shuffle=False,pin_memory=True)
+    """
+    set model seed 
+    """
+    Seed = args.man_seed if args.manualSeed else random.randint(1, 10000)
+    Seed=3407
+    print("Random Seed: ", Seed)
+    print(args)
+    random.seed(Seed)
+    torch.manual_seed(Seed)  
+    np.random.seed(Seed)    
+    main(args,train_loader,val_loader,test_loader)
+    

util.py

@@ -0,0 +1,102 @@
+import torch
+import pandas as pd
+import numpy as np
+
+from matplotlib import cm
+import matplotlib.pyplot as plt
+import scipy
+import torch.nn.functional as F
+import torchvision
+
+from sklearn.metrics import explained_variance_score,mean_squared_error,mean_absolute_error,r2_score,precision_score,recall_score,f1_score,roc_auc_score,roc_curve, auc,confusion_matrix
+from sklearn.feature_selection import r_regression
+
+from torch_sparse import SparseTensor
+from scipy.sparse import csr_matrix
+from scipy.sparse.csgraph import minimum_spanning_tree
+from math import pi as PI
+
+def scipy_spanning_tree(edge_index, edge_weight,num_nodes ):
+    row, col = edge_index.cpu()
+    edge_weight=edge_weight.cpu()
+    cgraph = csr_matrix((edge_weight, (row, col)), shape=(num_nodes, num_nodes))
+    Tcsr = minimum_spanning_tree(cgraph)
+    tree_row, tree_col = Tcsr.nonzero()
+    spanning_edges = np.stack([tree_row,tree_col],0)    
+    return spanning_edges
+    
+def build_spanning_tree_edge(edge_index,edge_weight, num_nodes):
+    spanning_edges = scipy_spanning_tree(edge_index, edge_weight,num_nodes,)
+        
+    spanning_edges = torch.tensor(spanning_edges, dtype=torch.long, device=edge_index.device)
+    spanning_edges_undirected = torch.cat([spanning_edges,torch.stack([spanning_edges[1],spanning_edges[0]])],1)
+    return spanning_edges_undirected
+
+
+
+
+def record(values,epoch,writer,phase="Train"):
+    """ tfboard write """
+    for key,value in values.items():
+        writer.add_scalar(key+"/"+phase,value,epoch)           
+def calculate(y_hat,y_true,y_hat_logit):
+    """ calculate five metrics using y_hat, y_true, y_hat_logit """
+    train_acc=(np.array(y_hat) == np.array(y_true)).sum()/len(y_true) 
+    # recall=recall_score(y_true, y_hat,zero_division=0,average='micro')
+    # precision=precision_score(y_true, y_hat,zero_division=0,average='micro')
+    # Fscore=f1_score(y_true, y_hat,zero_division=0,average='micro')
+    # roc=roc_auc_score(y_true, scipy.special.softmax(np.array(y_hat_logit),axis=1)[:,1],average='micro',multi_class='ovr')
+    # one_hot_encoded_labels = np.zeros((len(y_true), 100))
+    # one_hot_encoded_labels[np.arange(len(y_true)), y_true] = 1
+    # roc=roc_auc_score(one_hot_encoded_labels, scipy.special.softmax(np.array(y_hat_logit),axis=1),average='micro',multi_class='ovr')
+    return train_acc
+
+
+def print_1(epoch,phase,values,color=None):
+    """ print epoch info"""
+    if color is not None:
+        print(color( f"epoch[{epoch:d}] {phase}"+ " ".join([f"{key}={value:.3f}" for key, value in values.items()]) ))
+    else:
+        print(( f"epoch[{epoch:d}] {phase}"+ " ".join([f"{key}={value:.3f}" for key, value in values.items()]) ))
+
+def get_angle(v1, v2):
+    if v1.shape[1]==2:
+        v1=F.pad(v1, (0, 1),value=0)
+    if v2.shape[1]==2:
+        v2= F.pad(v2, (0, 1),value=0)
+    return torch.atan2( torch.cross(v1, v2, dim=1).norm(p=2, dim=1), (v1 * v2).sum(dim=1))
+def get_theta(v1, v2):
+    # v1 is starting line, right-hand rule to v2, if thumb is up, +, else -
+    angle=get_angle(v1, v2)
+    if v1.shape[1]==2:
+        v1=F.pad(v1, (0, 1),value=0)
+    if v2.shape[1]==2:
+        v2= F.pad(v2, (0, 1),value=0)
+    v = torch.cross(v1, v2, dim=1)[...,2]
+    flag = torch.sign((v))
+    flag[flag==0]=-1 
+    return angle*flag   
+
+def triplets(edge_index, num_nodes):
+    row, col = edge_index
+
+    value = torch.arange(row.size(0), device=row.device)
+    adj_t = SparseTensor(row=row, col=col, value=value,
+                         sparse_sizes=(num_nodes, num_nodes))
+    adj_t_col = adj_t[:,row]
+    num_triplets = adj_t_col.set_value(None).sum(dim=0).to(torch.long)
+
+    idx_j = row.repeat_interleave(num_triplets) 
+    idx_i = col.repeat_interleave(num_triplets) 
+    edx_2nd = value.repeat_interleave(num_triplets) 
+    idx_k = adj_t_col.t().storage.col() 
+    edx_1st = adj_t_col.t().storage.value()
+    mask1 = (idx_i == idx_k) & (idx_j != idx_i)  # Remove go back triplets. 
+    mask2 = (idx_i == idx_j) & (idx_j != idx_k)  # Remove repeat self loop triplets
+    mask3 = (idx_j == idx_k) & (idx_i != idx_k)  # Remove self-loop neighbors 
+    mask = ~(mask1 | mask2 | mask3) 
+    idx_i, idx_j, idx_k, edx_1st, edx_2nd = idx_i[mask], idx_j[mask], idx_k[mask], edx_1st[mask], edx_2nd[mask]
+    
+    num_triplets_real = torch.cumsum(num_triplets, dim=0) - torch.cumsum(~mask, dim=0)[torch.cumsum(num_triplets, dim=0)-1]
+
+    return torch.stack([idx_i, idx_j, idx_k]), num_triplets_real.to(torch.long), edx_1st, edx_2nd