graph_classification_param_sens.py

import socket
import warnings
import torch
import os
import os.path as osp
import GCL.losses as L
import GCL.augmentors as A
import torch.nn.functional as F
import GCL.augmentors.learnable_augs as LeA
import GCL.augmentors.manually_augs as MaA
import logging
import joblib
import argparse
import numpy as np
import networkx as nx
import sys
import time

from torch import nn
from tqdm import tqdm
from torch.optim import Adam
from GCL.eval import get_split, SVMEvaluator
from GCL.models import DualBranchContrast
from torch_geometric.nn import GINConv, global_add_pool
from torch_geometric.data import DataLoader
from GCL.models.mymodels import GraphNodeEncoder, GraphEncoder, Predictor
from GCL.utils import (compute_infonce, cluster_get, CustomDataLoader, compute_cluster_constrain_loss, 
                       k_fold, save_model_evaluator, load_model_evaluator, str_to_bool, add_extra_pos_mask, 
                       degree, topo_cluster_labels_get, compute_spectral_topo_loss, compute_spectral_feature_loss, 
                       print_memory_usage, print_args)
from sklearn.metrics import f1_score
from datetime import datetime
from general_data_loader import load_dataset_graphcls, get_split_mask, CombinedDataset, get_subgraph_edge_attr_with_postfix
from torch_geometric.utils import scatter, to_networkx, to_dense_adj, get_laplacian
from sklearn.model_selection import KFold


def get_config():
    parser = argparse.ArgumentParser()
    
    parser.add_argument("--load_subgraph", type=str_to_bool, default=True)
    parser.add_argument("--use_cluster_reg", type=str_to_bool, default=False)
    parser.add_argument("--use_degree", type=str_to_bool, default=False)
    parser.add_argument("--use_spectral_fea", type=str_to_bool, default=True)
    parser.add_argument("--use_spectral_topo", type=str_to_bool, default=True)
    parser.add_argument("--linear", type=str_to_bool, default=False)

    parser.add_argument("--cluster_reg_weight", type=float, default=0.1)
    parser.add_argument("--spectral_topo_weight", type=float, default=0.8) # 0.8
    parser.add_argument("--spectral_feature_weight", type=float, default=1.0) # 0.3
    parser.add_argument("--sim_weight", type=float, default=0.0) # 0.1

    parser.add_argument("--sample_edges_ratio", type=float, default=1.0) # topo augmentation
    parser.add_argument("--ratio_of_edge_insert", type=float, default=0.0)
    parser.add_argument("--control_by_ratio", type=str_to_bool, default=False)
    parser.add_argument("--drop_edge_ratio", type=float, default=0.0)
    parser.add_argument("--num_drop_ratio", type=float, default=0.1) # feature augmentation
    

    parser.add_argument("--mode", type=str, default="param_sens", choices=["unsup", "semisup", "param_sens"])
    parser.add_argument("--dataset_name", type=str, default="MUTAG")
    parser.add_argument("--epoch_select", type=str, default="val_max")


    parser.add_argument("--gpu", type=int, default=0)
    parser.add_argument("--batch_size", type=int, default=512) # 512
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--n_clusters", type=int, default=5)
    parser.add_argument("--par_id", type=int, default=0)

    args = parser.parse_args()
    return args


def train(encoder_model, predictor, dataloader, optimizer, args, augs_type=None, device='cpu'):
    encoder_model.train()
    epoch_loss = 0
    epoch_cluster_reg = 0
    epoch_predict_loss = 0
    epoch_spectral_topo_loss = 0
    epoch_spectral_feature_loss = 0
    encoder_exec_time = 0
    spectral_exec_time = 0
    epoch_sim_loss = 0
    epoch_fea_sim_loss = 0
    for data, extra_info_dict in dataloader:
        """
        data.cluster_labels: list of tensor: [n_nodes, 1]
        data.cluster_centrics: list of tensor: [n_cluster, fea_dim]
        data.cluster_num: list of int
        """
        data = data.to(device)

        # cluster_centrics = [c.to(device) for c in cluster_centrics]
        # cluster_labels = [cl.to(device) for cl in cluster_labels]
        # topo_cluster_labels = [cl.to(device) for cl in topo_cluster_labels]
        eigen_values = extra_info_dict["eigen_values"]
        eigen_values = [ev.to(device) for ev in eigen_values]
        if "x_eigen_values" in extra_info_dict:
            x_eigen_values = extra_info_dict["x_eigen_values"]
            x_eigen_vectors_U = extra_info_dict["x_eigen_vectors_U"]
            x_eigen_vectors_V = extra_info_dict["x_eigen_vectors_V"]

            x_eigen_values = [ev.to(device) for ev in x_eigen_values]
            x_eigen_vectors_U = [elem.to(device) for elem in x_eigen_vectors_U]
            x_eigen_vectors_V = [elem.to(device) for elem in x_eigen_vectors_V]

        else:
            x_eigen_values = None
            x_eigen_vectors_U = None
            x_eigen_vectors_V = None
        
        # eigen_vectors = [ev.to(device) for ev in eigen_vectors]
        assert len(eigen_values)==(data.batch.max()+1)
        assert data.num_component.shape[0]==(data.batch.max()+1)
        assert data.num_nodes==data.eigen_vectors.shape[0]
        # make sure edge_weights [num_edges, 1]
        if 'edge_weights' in data.keys():
            edge_weights = data.edge_weights.reshape(-1,1)
        else:
            edge_weights = torch.ones((data.edge_index.shape[1], 1), device=device)

        batch_size = len(data)
        num_nodes = data.batch.size(0)
        optimizer.zero_grad()

        if args.use_degree:
            if data.x is None:
                data.x = data.degree
            # else:
            #     data.x = torch.concat([data.x, data.degree], dim=1)
        else:
            if data.x is None:
                data.x = torch.ones((num_nodes, 1), dtype=torch.float32, device=data.batch.device)
        
        # edge attr construct
        edge_attr = torch.pow(data.eigen_vectors[data.edge_index[0]] - data.eigen_vectors[data.edge_index[1]], 2)
        edge_attr = torch.cat([edge_attr, torch.concat([data.x[data.edge_index[0]], data.x[data.edge_index[1]]], dim=1)], dim=1)

        if x_eigen_values!=None:
            all_x_eigen_distance = []
            for i in range(len(x_eigen_vectors_U)): # for each single graph
                U_expanded = x_eigen_vectors_U[i].unsqueeze(1) # [n_nodes, 1, n_eigen_vecs]
                V_expanded = x_eigen_vectors_V[i].unsqueeze(0) # [1, n_feas, n_eigen_vecs]

                all_x_eigen_distance.append(torch.pow(U_expanded - V_expanded, 2)) # [n_nodes, n_feas, n_eigen_vecs]
            x_eigen_distance = torch.concat(all_x_eigen_distance)
            # print(U_expanded.shape, V_expanded.shape, x_eigen_distance.shape, data.num_nodes, data.x.shape)
        else:
            x_eigen_distance = None
        
        # print(encoder_model.device, encoder_model.aug2.device)
        start_time = time.time()
        _, g, _, g1, _, g2, augs_g = encoder_model(data.x, data.edge_index, data.batch, edge_weights=edge_weights, edge_attr=edge_attr, 
                                                   eigen_values=eigen_values, eigen_vectors=data.eigen_vectors, x_eigen_distance=x_eigen_distance)
        g1, g2 = [encoder_model.graph_node_encoder.project(g) for g in [g1, g2]]
        # loss = contrast_model(g1=g1, g2=g2, batch=data.batch)
        encoder_exec_time += time.time() - start_time

        # denote positive samples, only for semi-supervise
        trues = data.y
        labels = data.y.unique()
        num_samples = data.y.shape[0]
        pos_mask = torch.eye(num_samples, dtype=bool, device=device)
        if args.mode == "semisup":
            pos_mask = add_extra_pos_mask(pos_mask, data, data.train_mask)
        # for label in labels:
        #     label_idx_vector = torch.zeros((num_samples, 1), dtype=bool, device=device)
        #     indxs_for_label = torch.argwhere(trues==label).squeeze()
        #     label_idx_vector[indxs_for_label] = True
        #     pos_mask += label_idx_vector.mul(label_idx_vector.T)
        
        if args.use_cluster_reg:
            aug_1, aug_2 = augs_g
            aug_x1 = aug_1[0]
            aug_x2 = aug_2[0]
            for b in range(batch_size):
                if args.use_degree:
                    single_graph_x = aug_x2[data.batch==b][:,:-1] # get node attr, last dim is degree feature
                else:
                    single_graph_x = aug_x2[data.batch==b]
                # cluster_reg = compute_cluster_constrain_loss(single_graph_x, cluster_labels[b], cluster_centrics[b], cluster_num[b], device=device)
            
            epoch_cluster_reg += cluster_reg.item()
        else:
            cluster_reg = 0
        
        ssl_loss = compute_infonce(g1, g2, pos_mask)
        epoch_loss += ssl_loss.item()

        if args.mode == "semisup":
            trues = data.y[data.train_mask]
            preds1 = predictor(g1)[data.train_mask]
            preds2 = predictor(g2)[data.train_mask]
            predict_loss1 = F.nll_loss(preds1, trues, reduction='sum')/(trues.shape[0])
            predict_loss2 = F.nll_loss(preds2, trues, reduction='sum')/(trues.shape[0])
            predict_loss = (predict_loss1 + predict_loss2)/2
            epoch_predict_loss += predict_loss.item()
        else:
            predict_loss = 0

        start_time = time.time()
        # spectral topological regularization
        aug_g1, aug_g2 = augs_g
        spectral_topo_loss = 0
        if args.use_spectral_topo:
            if augs_type[0] in ["LearnableEdgeDropping", "LearnableEdgeAdding", "LearnableEdgePerturbation"]:
                spectral_topo_loss += compute_spectral_topo_loss(data.edge_index, edge_weights, aug_g1, eigen_values, data.num_component, data.batch, augs_type[0])
                epoch_spectral_topo_loss += spectral_topo_loss.item()
            if augs_type[1] in ["LearnableEdgeDropping", "LearnableEdgeAdding", "LearnableEdgePerturbation"]:
                spectral_topo_loss += compute_spectral_topo_loss(data.edge_index, edge_weights, aug_g2, eigen_values, data.num_component, data.batch, augs_type[1])
                epoch_spectral_topo_loss += spectral_topo_loss.item()
            spectral_topo_loss = spectral_topo_loss/len(data)
            

        # spectral feature regularization
        spectral_feature_loss = 0
        if args.use_spectral_fea:
            if x_eigen_values!=None:
                if augs_type[0] in ["LearnableFeatureDroppingBySpectral"]:
                    spectral_feature_loss += compute_spectral_feature_loss(aug_g1, x_eigen_values, data.batch)/(data.batch.max()+1)
                    epoch_spectral_feature_loss += spectral_feature_loss.item()
                if augs_type[1] in ["LearnableFeatureDroppingBySpectral"]:
                    spectral_feature_loss += compute_spectral_feature_loss(aug_g2, x_eigen_values, data.batch)/(data.batch.max()+1)
                    epoch_spectral_feature_loss += spectral_feature_loss.item()
        spectral_exec_time += time.time() - start_time

        # sim loss
        sim_loss = 0
        # add_edge_weights = aug_g2[6]
        # sim_loss = -(torch.norm(add_edge_weights)**2)/len(data)
        # epoch_sim_loss += sim_loss.item()
        # auged_x = aug_g1[0]
        # fea_sim_loss = -(torch.norm(auged_x-data.x)**2)/len(data)
        # epoch_fea_sim_loss += fea_sim_loss.item()

        loss = ssl_loss + predict_loss + args.spectral_topo_weight*spectral_topo_loss + args.spectral_feature_weight*spectral_feature_loss + args.sim_weight*sim_loss

        loss.backward() # on


        optimizer.step()
    # print_memory_usage()
    # print("encoder running time: ", encoder_exec_time)
    # print("spectral running time: ", spectral_exec_time)

    return epoch_loss, epoch_predict_loss, epoch_spectral_topo_loss, epoch_spectral_feature_loss, epoch_sim_loss, epoch_fea_sim_loss


def test(encoder_model, dataloader, args, split, epoch_select="val_max", best_evaluator=None, device='cpu'):
    encoder_model.eval()
    x = []
    y = []
    for data, _ in dataloader:
        data = data.to(device)
        num_nodes = data.batch.size(0)
        if args.use_degree:
            if data.x is None:
                data.x = data.degree
            # else:
            #     data.x = torch.concat([data.x, data.degree], dim=1)
        else:
            if data.x is None:
                data.x = torch.ones((num_nodes, 1), dtype=torch.float32, device=data.batch.device)
        
        if 'edge_weights' in data.keys():
            edge_weights = data.edge_weights.reshape(-1,1)
        else:
            edge_weights = torch.ones((data.edge_index.shape[1], 1), device=device)
        
        z, g = encoder_model(data.x, data.edge_index, data.batch, edge_weights=edge_weights, mode="eval")
        x.append(g)
        y.append(data.y)
    x = torch.cat(x, dim=0)
    y = torch.cat(y, dim=0)

    # train_mask, val_mask, test_mask = get_split_mask(x.shape[0], train_ratio=0.8, test_ratio=0.1)
    # split = {"train": train_mask,
    #          "valid": val_mask,
    #          "test": test_mask}
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        if best_evaluator==None:
            result, best_evaluator = SVMEvaluator(linear=args.linear, epoch_select=epoch_select, dataset_name=args.dataset_name)(x, y, split)
        else:
            x_test = x[split['test']].detach().cpu().numpy()
            y_test = y[split['test']].detach().cpu().numpy()
            test_micro = f1_score(y_test, best_evaluator.predict(x_test), average='micro')
            test_macro = f1_score(y_test, best_evaluator.predict(x_test), average='macro')
            result = {
                'micro_f1': test_micro,
                'macro_f1': test_macro,
            }
    return result, best_evaluator



def main():
    hostname = socket.gethostname()
    print("Current hostname:", hostname)

    EIGEN_VEC_NUM = 6 # 50 4
    load_dataset = True
    x_eigen_load_dataset = True
    preprocess_save = True
    not_load_exist_model = False
    val_gap=20
    save_checkpoint_gap=50
    args = get_config()
    if args.dataset_name=="COLLAB":
        args.batch_size = 512-80
        args.epochs = 300
        val_gap=10
    elif args.dataset_name=="REDDIT-BINARY":
        args.batch_size = 512-160
    elif args.dataset_name=="REDDIT-MULTI-5K":
        args.batch_size = 256
    elif args.dataset_name=="DD":
        args.batch_size = 256
    elif args.dataset_name=="NCI1":
        args.batch_size = 512
        args.epochs = 300
    

    batch_size = args.batch_size
    epochs = args.epochs
    dataset_name = args.dataset_name

    if torch.cuda.is_available():
        device = torch.device(f"cuda:{args.gpu}")
    else:
        device = torch.device("cpu")
    # device = torch.device("cpu")
        

    current_time = datetime.now()
    print("Current time:", current_time)
    # print settings
    print_args(args)
    print("dataset: ", args.dataset_name, "device: ", device, ", mode: ", args.mode, ", cluster_reg: ", args.use_cluster_reg, 
        ", spectral_topo_weight: ", args.spectral_topo_weight, ", spectral_feature_weight: ", args.spectral_feature_weight, 
        ", is linear: ", args.linear, ", EIGEN_VEC_NUM: ", EIGEN_VEC_NUM)
    
    if not osp.exists('./log'):
        os.mkdir('./log')
    if not osp.exists(f'./log/{dataset_name}'):
        os.mkdir(f'./log/{dataset_name}')
    logging_path = "./log/{}/results_{}.log".format(dataset_name, args.mode)
    logging.basicConfig(filename=logging_path, level=logging.DEBUG, format='%(asctime)s %(message)s')

    if not osp.exists('./checkpoints'):
        os.mkdir('./checkpoints')
    if not osp.exists(f'./checkpoints/{dataset_name}'):
        os.mkdir(f'./checkpoints/{dataset_name}')
    if not osp.exists(f'./checkpoints/{dataset_name}/{args.mode}'):
        os.mkdir(f'./checkpoints/{dataset_name}/{args.mode}')
    # checkpoints_path = "./checkpoints/{}".format(dataset_name)
    
    # data load
    path = osp.join(osp.expanduser('~'), 'datasets')
    dataset, split = load_dataset_graphcls(path, dataset_name, args)

    # description
    data_summary = dataset.get_summary()
    num_nodes = dataset.edge_index.max()+1
    num_node_features = dataset.num_node_features
    num_classes = dataset.num_classes
    print(data_summary)
    print(dataset_name, "num_nodes: ", num_nodes, "num_node_features: ", num_node_features, "num_classes: ", num_classes)

    if num_node_features>=1:
        x_norm_part = torch.linalg.norm(dataset.x, dim=0).to(device)
        x_norm_part = torch.clip(x_norm_part, min=1e-6)
    else:
        x_norm_part = 1
    

    dataset_dict = {"dataset": dataset}
    # for EIGEN_VEC_NUM in [2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40]:
    #     print(f"###################### Parameter EIGEN_VEC_NUM: {EIGEN_VEC_NUM} ############################")
    
    # for topo_weight in [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1, 2, 3, 4, 5, 10, 20, 30, 50, 100]:
    #     args.spectral_topo_weight = topo_weight
    #     print(f"###################### Parameter topo_weight: {args.spectral_topo_weight} ############################")

    # for feature_weight in [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1, 2, 3, 4, 5, 10, 20, 30, 50, 100]:
    #     args.spectral_feature_weight = feature_weight
    #     print(f"###################### Parameter feature_weight: {args.spectral_feature_weight} ############################")

    # for pert_ratio in [0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5]:
    #     args.sample_edges_ratio = 1.0
    #     args.ratio_of_edge_insert = pert_ratio
    #     args.control_by_ratio = True
    #     args.drop_edge_ratio = pert_ratio
    #     args.num_drop_ratio = pert_ratio
    #     print(f"###################### Parameter pert_ratio: {pert_ratio} ############################")
        
    for trial in range(10):
        print(f"###################### Trail: {trial} ############################")



        # degree calculating
        if EIGEN_VEC_NUM==50:
            whole_processed_dataset_save_path = osp.join(path, dataset_name, "whole_processed_dataset.pt")
        else:
            whole_processed_dataset_save_path = osp.join(path, dataset_name, f"whole_processed_dataset_{EIGEN_VEC_NUM}.pt")
        if osp.exists(whole_processed_dataset_save_path) and load_dataset:
            id_degree_dict = torch.load(whole_processed_dataset_save_path)
        else:
            all_degrees = []
            all_edge_weights = []
            all_graph_id = []
            all_eigen_values = [] # all non-zero eigen values
            all_eigen_vectors = [] # and corresponded eigen vectors
            with tqdm(total=len(dataset), desc='(PRE)') as pbar:
                for i, data in enumerate(dataset):
                    node_degree = degree(data.edge_index)
                    graph_degree = node_degree.sum()
                    node_degree_features = (node_degree/graph_degree).reshape(-1, 1)
                    assert node_degree_features.shape[0] == data.num_nodes
                    all_degrees.append(node_degree_features)
                    pbar.update()

                    edge_weights = torch.ones((data.edge_index.shape[1],))
                    all_edge_weights.append(edge_weights)

                    all_graph_id.append(i)
            # combine dataset
            id_degree_dict = {"graph_id": all_graph_id, 
                              "degree": all_degrees, 
                              "edge_weights":all_edge_weights, 
                            }
            if preprocess_save:
                torch.save(id_degree_dict, whole_processed_dataset_save_path)
        
        dataset_dict.update({'graph_id': id_degree_dict["graph_id"]})
        dataset_dict.update({'degree': id_degree_dict["degree"]})
        dataset_dict.update({'train_mask': torch.zeros(len(dataset), dtype=bool)})
        dataset_dict.update({'edge_weights': id_degree_dict['edge_weights']})


        # eigen info calculate
        if EIGEN_VEC_NUM==50:
            extra_eig_info_dataset_save_path = osp.join(path, dataset_name, "eig_info.pt")
        else:
            extra_eig_info_dataset_save_path = osp.join(path, dataset_name, f"eig_info_{EIGEN_VEC_NUM}.pt")
        if osp.exists(extra_eig_info_dataset_save_path) and load_dataset:
            eig_info_dict = torch.load(extra_eig_info_dataset_save_path)
        else:
            all_eigen_values = [] # all non-zero eigen values
            all_eigen_vectors = [] # and corresponded eigen vectors
            all_num_components = []
            with tqdm(total=len(dataset), desc='(Eig)') as pbar:
                for i, data in enumerate(dataset):
                    g = to_networkx(data, to_undirected=True)
                    num_components = nx.number_connected_components(g)
                    all_num_components.append(num_components)
                    del g

                    # Normalized laplacian
                    lap_indicis, lap_values = get_laplacian(data.edge_index, normalization='sym')
                    norm_L = torch.sparse_coo_tensor(lap_indicis, lap_values, size=(data.num_nodes, data.num_nodes)).to_dense()

                    eig_vals, eig_vecs = torch.linalg.eig(norm_L)
                    indices = eig_vals.real.sort().indices[num_components:num_components+EIGEN_VEC_NUM] # ascending
                    all_eigen_values.append(eig_vals[indices].real)
                    eig_vecs = eig_vecs[:, indices].real
                    if eig_vecs.shape[1] < EIGEN_VEC_NUM:
                        zero_columns = torch.zeros(eig_vecs.shape[0], EIGEN_VEC_NUM - eig_vecs.shape[1])
                        eig_vecs = torch.cat((eig_vecs, zero_columns), dim=1)
                    all_eigen_vectors.append(eig_vecs)
                    
                    pbar.update()

            eig_info_dict = {"eigen_values": all_eigen_values, 
                            "eigen_vectors": all_eigen_vectors, 
                            "num_component": all_num_components,
                            }
            if preprocess_save:
                torch.save(eig_info_dict, extra_eig_info_dataset_save_path)
        dataset_dict.update({'num_component': eig_info_dict["num_component"]})
        dataset_dict.update({'eigen_values': eig_info_dict["eigen_values"]})
        dataset_dict.update({'eigen_vectors': eig_info_dict["eigen_vectors"]})


        if num_node_features >= 1:
            if EIGEN_VEC_NUM==50:
                extra_x_eig_info_dataset_save_path = osp.join(path, dataset_name, "x_eig_info.pt")
            else:
                extra_x_eig_info_dataset_save_path = osp.join(path, dataset_name, f"x_eig_info_{EIGEN_VEC_NUM}.pt")
            X_EIGEN_VEC_NUM = min(num_node_features, EIGEN_VEC_NUM)
            if osp.exists(extra_x_eig_info_dataset_save_path) and x_eigen_load_dataset:
                x_eig_info_dict = torch.load(extra_x_eig_info_dataset_save_path)
            else:
                # feature x's sigular vector: U, S, V
                all_x_eigen_values = [] # all non-zero eigen values
                all_x_eigen_vectors_U = []
                all_x_eigen_vectors_V = [] # and corresponded eigen vectors
                with tqdm(total=len(dataset), desc='(X_Eig)') as pbar:
                    for i, data in enumerate(dataset):
                        x_norm = data.x/x_norm_part.cpu()
                        U, S, Vh = torch.linalg.svd(x_norm, full_matrices=False)
                        U = U[:, :X_EIGEN_VEC_NUM]
                        V = Vh.T[:, :X_EIGEN_VEC_NUM]
                        all_x_eigen_values.append(S[:X_EIGEN_VEC_NUM]) # descending

                        # expand_X = F.pad(data.x, (data.x.shape[0], 0, 0, data.x.shape[1]), mode='constant', value=0)
                        # A_X = expand_X+expand_X.T
                        # eig_vals, eig_vecs = torch.linalg.eig(A_X)
                        # indices = eig_vals.real.sort().indices[1:1+X_EIGEN_VEC_NUM]
                        # all_x_eigen_values.append(eig_vals[indices].real)
                        # U = eig_vecs[:data.x.shape[0], indices].real
                        # V = eig_vecs[data.x.shape[0]:, indices].real

                        if U.shape[1] < X_EIGEN_VEC_NUM:
                            zero_columns = torch.zeros(U.shape[0], X_EIGEN_VEC_NUM - U.shape[1])
                            U = torch.cat((U, zero_columns), dim=1)
                        if V.shape[1] < X_EIGEN_VEC_NUM:
                            zero_columns = torch.zeros(V.shape[0], X_EIGEN_VEC_NUM - V.shape[1])
                            V = torch.cat((V, zero_columns), dim=1)

                        all_x_eigen_vectors_U.append(U)
                        all_x_eigen_vectors_V.append(V)

                        pbar.update()
                x_eig_info_dict = {"x_eigen_values": all_x_eigen_values,
                                "x_eigen_vectors_U": all_x_eigen_vectors_U, 
                                "x_eigen_vectors_V": all_x_eigen_vectors_V,
                                }
                if preprocess_save:
                    torch.save(x_eig_info_dict, extra_x_eig_info_dataset_save_path)
            dataset_dict.update({'x_eigen_values': x_eig_info_dict["x_eigen_values"]})
            dataset_dict.update({'x_eigen_vectors_U': x_eig_info_dict["x_eigen_vectors_U"]})
            dataset_dict.update({'x_eigen_vectors_V': x_eig_info_dict["x_eigen_vectors_V"]})


        if (num_node_features>=1):
            dataset_dict_keys = ["dataset", "graph_id", "train_mask", "degree", "edge_weights", 
                                'num_component', "eigen_values", "eigen_vectors", 
                                "x_eigen_values", "x_eigen_vectors_U", "x_eigen_vectors_V"]
        else:
            dataset_dict_keys = ["dataset", "graph_id", "train_mask", "degree", "edge_weights", 
                                'num_component', "eigen_values", "eigen_vectors"]
        dataset_dict = {key: dataset_dict[key] for key in dataset_dict_keys}

        # # topo logical clustering
        # topo_cluster_info_save_path = osp.join(path, dataset_name, "topo_cluster_infos.pt")
        # if osp.exists(topo_cluster_info_save_path):
        #     topo_cluster_info_dict = torch.load(topo_cluster_info_save_path)
        #     topo_cluster_labels_list = topo_cluster_info_dict['cluster_labels_list']
        # else:
        #     topo_cluster_labels_list = topo_cluster_labels_get(dataset, n_clusters)
        #     topo_cluster_infos = {"cluster_labels_list": topo_cluster_labels_list}
        #     torch.save(topo_cluster_infos, topo_cluster_info_save_path)
        #     print("topo cluster_infos: ", n_clusters)

        # dataset_dict.update({'topo_cluster_labels_list': topo_cluster_labels_list})

        combined_dataset = CombinedDataset(**dataset_dict)
        # combined_dataset = list(zip(dataset, cluster_labels_list, cluster_centrics_list, cluster_num_list, topo_cluster_labels_list, split['train']))
        print("combined_dataset: ", type(combined_dataset[0][0]), type(combined_dataset[0][1]), type(combined_dataset[0][2]), type(combined_dataset[0][3]))

        # create custom dataloader
        # dataloader = DataLoader(dataset, batch_size=batch_size)
        dataloader = CustomDataLoader(dataset=combined_dataset, batch_size=batch_size, shuffle=True, has_node_features=(num_node_features>=1))


        if args.use_degree:
            input_fea_dim = max(num_node_features, 1)
        else:
            input_fea_dim = max(num_node_features, 1)

        # define augmentations
        aug1 = A.Identity()
        if num_node_features >= 1:
            x_eigen_distance_input_dim = X_EIGEN_VEC_NUM
            leA_FD = LeA.LearnableFeatureDroppingBySpectral(input_dim=x_eigen_distance_input_dim, hidden_dim=128, num_drop_ratio=args.num_drop_ratio).to(device)
            aug1 = leA_FD
        # aug1 = MaA.CrossClusterEdgeDropping(ratio=0.1)
        # aug2 = A.Identity()
        rand_FM = A.FeatureMasking(pf=0.1)
        rand_EA = A.EdgeAdding(pe=0.3)
        rand_ED = A.EdgeRemoving(pe=0.3)
        feadrop_input_dim = int(data_summary.num_nodes.quantile25 * batch_size)//2
        # print("feadrop_input_dim: ", feadrop_input_dim)
        # aug2 = LeA.LearnableFeatureDropping(input_dim=feadrop_input_dim, hidden_dim=16, temp=1.0).to(device)
        # if dataset.num_node_features > 1:
        #     x_eigen_distance_input_dim = X_EIGEN_VEC_NUM
        #     leA_FD = LeA.LearnableFeatureDroppingBySpectral(input_dim=x_eigen_distance_input_dim, hidden_dim=128).to(device)
        #     aug2 = leA_FD

        edge_attr_input_dim = input_fea_dim*2 + dataset_dict['eigen_vectors'][0].shape[1]
        leA_ED = LeA.LearnableEdgeDropping(input_dim=edge_attr_input_dim, hidden_dim=128, temp=1.0).to(device) # edge_attr: subg feas, eigen_vecs
        leA_EA = LeA.LearnableEdgeAdding(input_dim=edge_attr_input_dim, hidden_dim=128, sample_edges_ratio=0.2).to(device)
        leA_EP = LeA.LearnableEdgePerturbation(input_dim_drop=edge_attr_input_dim, input_dim_add=edge_attr_input_dim, hidden_dim=128, 
                                               sample_edges_ratio=args.sample_edges_ratio, 
                                               ratio_of_edge_insert=args.ratio_of_edge_insert, 
                                               control_by_ratio=args.control_by_ratio, 
                                               drop_edge_ratio=args.drop_edge_ratio).to(device)
        # aug2 = leA_EP
        aug2 = A.Identity()
        # aug2 = LeA.Compose([leA_ED, leA_EA])
        # aug2 = MaA.CrossClusterEdgeAdding(ratio=0.3)
        # aug2 = MaA.WithinClusterEdgeDropping(ratio=0.3)
        # aug2 = MaA.WithinClusterEdgeAdding(ratio=0.3)
        # aug2 = MaA.CrossClusterEdgeDropping(ratio=0.3)
        augs_type = [type(aug1).__name__, type(aug2).__name__]
        aug1_no_spec = aug2_no_spec = False
        if augs_type[0] == "LearnableFeatureDroppingBySpectral":
            if args.use_spectral_fea==False:
                augs_type[0] += "WithoutSpectral"
                aug1_no_spec = True
                val_gap=epochs//2
        elif augs_type[0] in ["LearnableEdgeDropping", "LearnableEdgeAdding", "LearnableEdgePerturbation"]:
            if args.use_spectral_topo==False:
                augs_type[0] += "WithoutSpectral"
                aug1_no_spec = True
                val_gap=epochs
        
        if augs_type[1] == "LearnableFeatureDroppingBySpectral":
            if args.use_spectral_fea==False:
                augs_type[1] += "WithoutSpectral"
                aug2_no_spec = True
                val_gap=epochs//2
        elif augs_type[1] in ["LearnableEdgeDropping", "LearnableEdgeAdding", "LearnableEdgePerturbation"]:
            if args.use_spectral_topo==False:
                augs_type[1] += "WithoutSpectral"
                aug2_no_spec = True
                val_gap=epochs
        not_load_exist_model = (aug1_no_spec and aug2_no_spec)
        if not_load_exist_model:
            val_gap=epochs

        print(augs_type)
        # aug2 = (LeA.LearnableEdgeDropping())

        checkpoints_path = "./checkpoints/{}/{}/{}_{}".format(dataset_name, args.mode, augs_type[0], augs_type[1])
        if not osp.exists(checkpoints_path):
            os.mkdir(checkpoints_path)

        # define encoder
        gconv = GraphNodeEncoder(input_channels=input_fea_dim, hidden_channels=256, num_layers=2).to(device)
        encoder_model = GraphEncoder(graph_node_encoder=gconv, augmentor=(aug1, aug2), x_norm_part=x_norm_part).to(device)
        predictor = Predictor(gconv.project_dim, len(dataset.y.unique())).to(device) # num_classes = len(dataset.y.unique())
        
        # define optimizer
        params_for_optimizer = [{'params': encoder_model.parameters()}]
        if isinstance(aug1, LeA.Compose):
            for aug_elem in aug1.augmentors:
                if isinstance(aug_elem, LeA.LearnableAugmentor):
                    params_for_optimizer.append({"params": aug_elem.parameters()})
        elif isinstance(aug1, LeA.LearnableAugmentor):
            print(augs_type[0], "learnable")
            params_for_optimizer.append({"params": aug1.parameters()})
        if isinstance(aug2, LeA.Compose):
            for aug_elem in aug2.augmentors:
                if isinstance(aug_elem, LeA.LearnableAugmentor):
                    params_for_optimizer.append({"params": aug_elem.parameters()})
        elif isinstance(aug2, LeA.LearnableAugmentor):
            print(augs_type[1], "learnable")
            params_for_optimizer.append({"params": aug2.parameters()})
        optimizer = Adam(params_for_optimizer, lr=0.01)

        # training
        split = get_split(num_samples=len(dataset), train_ratio=0.8, test_ratio=0.1)
        best_acc = 0
        print("-"*40+f"Starting"+"-"*40)
        all_results = {'micro_f1': [], 
                    'macro_f1': []}
        all_max_results = {'micro_f1': [],
                        'macro_f1': []}
        all_min_results = {'micro_f1': [],
                        'macro_f1': []}
        test_acc = "nan"
        # print_memory_usage()
        with tqdm(total=epochs, desc='(T)', file=open("/dev/stdout", "a")) as pbar:
            for epoch in range(1, epochs+1):
                # print_memory_usage()
                loss, predict_loss, spectral_topo_loss, spectral_feature_loss, sim_loss, fea_sim_loss = train(encoder_model, predictor, dataloader, optimizer, args, augs_type=augs_type, device=device)
                
                if ((dataset_name!="COLLAB") or (dataset_name=="COLLAB" and epoch>150)) and epoch % val_gap == 0:
                    all_fold_results = {'micro_f1': [], 
                                        'macro_f1': []}
                    for fold_i, (train_idx, val_idx, test_idx) in enumerate(k_fold(len(dataset), epoch_select=args.epoch_select)):
                        split = {'train': torch.asarray(train_idx, device=device), 
                                'valid': torch.asarray(val_idx, device=device),
                                'test':  torch.asarray(test_idx, device=device)}
                        test_result, best_evaluator = test(encoder_model, dataloader, args, split, epoch_select=args.epoch_select, best_evaluator=None, device=device)

                        all_fold_results['micro_f1'].append(test_result['micro_f1'])
                        all_fold_results['macro_f1'].append(test_result['macro_f1'])
                    avg_results = {'micro_f1': np.mean(all_fold_results['micro_f1']), 
                                'macro_f1': np.mean(all_fold_results['macro_f1'])}
                    max_results = {'micro_f1': np.max(all_fold_results['micro_f1']), 
                                'macro_f1': np.max(all_fold_results['macro_f1'])}
                    min_results = {'micro_f1': np.min(all_fold_results['micro_f1']), 
                                'macro_f1': np.min(all_fold_results['macro_f1'])}
                    all_results['micro_f1'].append(avg_results['micro_f1'])
                    all_results['macro_f1'].append(avg_results['macro_f1'])
                    all_max_results['micro_f1'].append(max_results['micro_f1'])
                    all_max_results['macro_f1'].append(max_results['macro_f1'])
                    all_min_results['micro_f1'].append(min_results['micro_f1'])
                    all_min_results['macro_f1'].append(min_results['macro_f1'])
                    
                    test_acc = avg_results['micro_f1']
                    if test_acc > best_acc:
                        best_acc = test_acc
                        save_model_evaluator(checkpoints_path, encoder_model, best_evaluator, par_id=args.par_id)

                pbar.set_postfix({'loss': loss, 
                                'spec_topo': spectral_topo_loss, 
                                'spec_feas': spectral_feature_loss,
                                'sim_loss': sim_loss,
                                'fea_sim_loss': fea_sim_loss,
                                'val_acc': test_acc})
                pbar.update()

        # for epoch in range(epochs):
        #     loss = train(encoder_model, dataloader, optimizer, device)
        #     print(loss)
        
        # testing
        best_encoder_model = encoder_model
        best_encoder_model, best_evaluator, best_predictor = load_model_evaluator(checkpoints_path, par_id=args.par_id)
        best_encoder_model = best_encoder_model.to(device)

        all_fold_results = {'micro_f1': [], 
                    'macro_f1': []}
        for fold_i, (train_idx, val_idx, test_idx) in enumerate(k_fold(len(dataset), epoch_select=args.epoch_select)):
            split = {'train': torch.asarray(train_idx, device=device), 
                    'valid': torch.asarray(val_idx, device=device),
                    'test':  torch.asarray(test_idx, device=device)}
            
            test_result, best_evaluator = test(best_encoder_model, dataloader, args, split, epoch_select="test_max", best_evaluator=None, device=device) # only for NCI1, PROTEINS
            all_fold_results['micro_f1'].append(test_result['micro_f1'])
            all_fold_results['macro_f1'].append(test_result['macro_f1'])
            print(f"{fold_i}: {test_result}")
        avg_results = {'micro_f1': np.mean(all_fold_results['micro_f1']), 
                    'macro_f1': np.mean(all_fold_results['macro_f1'])}
        max_results = {'micro_f1': np.max(all_fold_results['micro_f1']), 
                    'macro_f1': np.max(all_fold_results['macro_f1'])}
        min_results = {'micro_f1': np.min(all_fold_results['micro_f1']), 
                    'macro_f1': np.min(all_fold_results['macro_f1'])}
        print(f"last test avg results: {avg_results}")
        
        # all_results['micro_f1'].append(avg_results['micro_f1'])
        # all_results['macro_f1'].append(avg_results['macro_f1'])
        # all_max_results['micro_f1'].append(max_results['micro_f1'])
        # all_max_results['macro_f1'].append(max_results['macro_f1'])
        # all_min_results['micro_f1'].append(min_results['micro_f1'])
        # all_min_results['macro_f1'].append(min_results['macro_f1'])

        best_idx = np.argmax(all_results['micro_f1'])
        best_results = {'micro_f1': all_results['micro_f1'][best_idx], 
                        'macro_f1': all_results['macro_f1'][best_idx]}
        best_max_result = {'micro_f1': all_max_results['micro_f1'][best_idx], 
                        'macro_f1': all_max_results['macro_f1'][best_idx]}
        best_min_result = {'micro_f1': all_min_results['micro_f1'][best_idx], 
                        'macro_f1': all_min_results['macro_f1'][best_idx]}
        # results
        # print(f'(E): Avg all folds test F1Mi={avg_results["micro_f1"]:.4f}, F1Ma={avg_results["macro_f1"]:.4f}')
        print(f'(E): Best test F1Mi={best_results["micro_f1"]:.4f}, F1Ma={best_results["macro_f1"]:.4f}')
        # logging.info("EIGEN_VEC_NUM: {}, topo weight: {}, fea_weight: {}, avg_results: {}, max: {}, min: {}"\
        #             .format(EIGEN_VEC_NUM, args.spectral_topo_weight, args.spectral_feature_weight, best_results, best_max_result, best_min_result))
        # logging.info("pert_ratio: {}, fea_drop_ratio: {}, avg_results: {}, max: {}, min: {}"\
        #             .format(args.ratio_of_edge_insert, args.num_drop_ratio, best_results, best_max_result, best_min_result))
        logging.info("aug_types {}, avg_results: {}, max: {}, min: {}"\
                    .format(augs_type , best_results, best_max_result, best_min_result))
        print("-"*90)

if __name__ == '__main__':
    main()