utils_copt.py

import pandas as pd
import numpy as np
import os
import argparse
import torch
import math
import numpy.linalg as linalg
import matplotlib.pyplot as plt
import networkx as nx
import pickle
import warnings
import sklearn.metrics
warnings.filterwarnings('ignore')

import pdb




res_dir = 'results'
data_dir = 'data'
device = 'cuda' if torch.cuda.is_available() else 'cpu'


def parse_args():
    parser = argparse.ArgumentParser()
    
    parser.add_argument("--verbose", dest='verbose', action='store_const', default=False, const=True, help='Print out verbose info during optimization')
    parser.add_argument("--seed", dest='fix_seed', action='store_const', default=False, const=True, help='Fix seed for reproducibility and fair comparison.')
    parser.add_argument("--n_epochs", dest='n_epochs', type=int, default=300, help='Number of COPT iterations during training.')
    parser.add_argument("--hike", dest='lr_hike', action='store_const', default=False, const=True, help='Use learning rate hiking. This is recommended for most applications.')
    parser.add_argument("--hike_interval", dest='hike_interval', type=int, default=15, help='Number of iterations having *low loss* between lr hikes.')
    parser.add_argument("--fast", dest='fast', action='store_const', default=False, const=True, help='Use fast kernel, in particular multiscale laplacian')
    parser.add_argument("--early_stopping", action='store_true', default=False, help='stop training early if loss repeatedly reaches below some threshold, useful for sketching.')
    
    parser.add_argument("--dataset_type", type=str, default='synthetic', help='dataset type, can be "synthetic" or "real"') 
    parser.add_argument("--dataset_name", type=str, default='PROTEINS', help='dataset name, will be ignored if dataset_type is synthetic')  
    parser.add_argument("--m", default=30, type=int, help='Number of vertices in graph X, e.g. query graph, useful especially when testing')
    parser.add_argument("--n", default=30, type=int, help='Number of vertices in graph Y, e.g. dataset graph, useful especially when testing')  

    parser.add_argument("--sinkhorn_iter", type=int, default=10, help='Number of Sinkhorn scaling iterations during optimization')
    parser.add_argument("--grid_search", dest='grid_search', action='store_const', default=False, const=True, help='grid search for learning SVC classifier')
    parser.add_argument("--compress_fac", default=-1, type=int, help='Factor of compression, e.g. 2 means reduce to half as many vertices')
    parser.add_argument("--got_it", dest='st_it', type=int, default=5, help='Number of Sinkhorn iterations for GOT')
    parser.add_argument("--got_tau", dest='st_tau', type=float, default=1, help='Number of Sinkhorn iterations for GOT')
    parser.add_argument("--got_n_sample", dest='st_n_samples', type=int, default=10, help='Number of samples in stochastic sampling for GOT')
    parser.add_argument("--got_n_epochs", dest='st_epochs', type=int, default=1000, help='Number of Sinkhorn iterations for GOT')
    parser.add_argument("--got_lr", dest='st_lr', type=float, default=.5, help='Number of Sinkhorn iterations for GOT')
    parser.add_argument("--gw_alpha", dest='gw_alpha', type=float, default=.8, help='alpha parameter for GW')
    parser.add_argument("--gw_metric", dest='gw_metric', type=str, default='sqeuclidean', help='features metric for GW')    
    
    opt = parser.parse_args()    
    return opt

def create_graph_lap(n):
    """
    Create graph laplacian of given size.
    """
    g = nx.random_geometric_graph(n, .5)
    
    
    Lx = nx.laplacian_matrix(g, range(n))
    Lx = np.array(Lx.todense())
    Lx = np.array([[0, -.3, -.9],[-.3, 0, 0],[-.9, 0,0]])
    return Lx

def graph_to_lap(g):
    """
    Get Laplacian from nx graph g
    """
    if not isinstance(g, nx.Graph):
        g = nx.from_numpy_array(g)
    Lx = nx.laplacian_matrix(g).todense()  
    Lx = torch.from_numpy(Lx).to(dtype=torch_dtype)
    return Lx

def lap_to_graph(L):
    
    if isinstance(L, torch.Tensor):
        L = L.cpu().numpy()
    L = L.copy()
    np.fill_diagonal(L, 0)    
    return nx.from_numpy_array(-L)

def canonicalize_mx(mx):
    mx1 = mx.clone()
    diag = mx.diag()
    idx = diag.argsort(dim=0)
    n_mx = len(mx)
    for i in range(n_mx):
        mx[i] = mx1[idx[i]]
    mx1 = mx.clone()
    for i in range(n_mx):
        mx[:, i] = mx1[:, idx[i]]
    
    return mx 
    
def create_graph(n, gtype=None, seed=None, params={}):
    total_iter = 100
    cnt = 0
    while True:
        if gtype == 'block':
            
            if params['n_blocks'] == 3:
                m = n // 3
                g = nx.stochastic_block_model([m, m, n-2*m],[[0.98,0.01,.01],[0.01,0.98,.01],[0.01,.01,.98]], seed=seed)
            elif params['n_blocks'] == 4:
                m = n // 4
                g = nx.stochastic_block_model([m, m, m, n-3*m],[[.97,0.01,0.01,.01],[.01,0.97,0.01,.01],[.01,0.01,0.97,.01],[.01,0.01,0.01,.97] ], seed=seed)
            else:
                m = n // 2
                g = nx.stochastic_block_model([m, n-m],[[0.99,0.01],[0.01,0.99]], seed=seed)
        elif gtype == 'strogatz':
            g = nx.connected_watts_strogatz_graph(n, max(n//4, 3), p=.05, seed=seed)
        elif gtype == 'random_regular':
            
            d = max(n//8, 2)
            if n*d % 2 == 1:
                n += 1
            g = nx.random_regular_graph(d, n, seed=seed)
        elif gtype == 'binomial':
            
            prob = params['prob']
            g = nx.binomial_graph(n, prob, seed=seed)
        elif gtype == 'barabasi':
            d = max(4, n//6)
            g = nx.barabasi_albert_graph(n, d, seed=seed)
        elif gtype == 'powerlaw_tree':            
            g = nx.random_powerlaw_tree(n, gamma=3, tries=1300, seed=seed)
        elif gtype == 'caveman':
            n_cliques = params['n_cliques']
            clique_sz = params['clique_sz']
            assert n_cliques * clique_sz == n
            g = nx.connected_caveman_graph(n_cliques, clique_sz)
        
        
        
        elif gtype == 'random_geometric': 
            radius = params['radius']
            g = nx.random_geometric_graph(n, radius, seed=seed)
        elif gtype == 'barbell': 
            
            g = nx.barbell_graph(n//2, 1)
        elif gtype == 'ladder': 
            g = nx.ladder_graph(n)            
        elif gtype == 'grid': 
            g = nx.grid_graph([n,n])            
        elif gtype == 'hypercube': 
            g = nx.hypercube_graph(n)            
        elif gtype == 'pappus': 
            g = nx.pappus_graph()            
        elif gtype == 'star': 
            g = nx.star_graph(n)   
        elif gtype == 'cycle': 
            g = nx.cycle_graph(n)   
        elif gtype == 'wheel': 
            g = nx.wheel_graph(n)   
        elif gtype == 'lollipop': 
            g = nx.lollipop_graph(n//2, 1)   
        else:
            raise Exception('graph type not supported ', gtype)

        remove_isolates(g)
        cnt += 1
        if nx.is_connected(g) or cnt > total_iter:
            if cnt > total_iter:
                g = g.subgraph(sorted(nx.connected_components(g), key=len)[-1]).copy()
            break
    return g

def fetch_data(dataset_name):
    
    
    data = torch.load('{}_lap.pt'.format(dataset_name))
    return data['lap'], data['labels'], data['target']

def fetch_data_graphs(dataset_name):
    
    try:
        data = torch.load('data/{}_lap.pt'.format(dataset_name))
    except Exception:
        raise Exception('Dataset {} graph data not created yet. More data can be created using the generateData.py script as in README.'.format(dataset_name))
    graphs = []
    for g in data['lap']:
        graph = lap_to_graph(g)
        '''
        for i in range(len(g)):
            graph.add_node(i)
        pdb.set_trace()
        '''
        graphs.append(graph)    
            
    return graphs, data['labels'], np.array(data['target'])

def view_graph(L, soft_edge=False, labels=None, name=''):
    plt.clf()
    if isinstance(L, torch.Tensor):        
        L = L.cpu().numpy()
    L = L.copy()
    
    np.fill_diagonal(L, 0)
    L *= -1
    
    g = nx.from_numpy_array(L)
    fig = plt.figure()
    plt.axis('off')
    ax = plt.gca()
    
    
    layout = nx.spring_layout(g)
    nx.draw_networkx_nodes(g, layout, node_size=500, alpha=0.5, cmap=plt.cm.RdYlGn, node_color='r', ax=ax)
    if labels is None:        
        nx.draw_networkx_labels(g, layout, font_color='w', font_weight='bold', font_size=15, ax=ax)
    else:
        
        nx.draw_networkx_labels(g, layout, labels=labels, font_color='k', font_size=12, ax=ax) 
    if soft_edge:
        
        
        
        elarge = [(u, v) for (u, v, d) in g.edges(data=True) if d['weight'] > .5] 
        esmall = [(u, v) for (u, v, d) in g.edges(data=True) if d['weight'] <= .5 and d['weight'] > .19] 
        nx.draw_networkx_edges(g, layout, edgelist=elarge, width=3.5, ax=ax)
        nx.draw_networkx_edges(g, layout, edgelist=esmall, wifth=3, ax=ax)
    else:
        nx.draw_networkx_edges(g, layout, ax=ax)
    
    fig.savefig('data/view_graph_{}.jpg'.format(name))
    print('plot saved to {}'.format('data/view_graph_{}.jpg'.format(name)))
    plt.show()

def plot_confusion(tgt, pred, labels=None, name=''):
    """
    Input:
    tgt, pred: target and predicted classes.
    labels: node labels.
    """
    plt.clf()
    fig = plt.figure()
    
    
    ax = plt.gca()
    mx = sklearn.metrics.confusion_matrix(tgt, pred)
    
    img = plt.matshow(mx)
    path = 'data/confusion_mx_{}.jpg'.format(name)
    
    ax.legend()
    name2label = {'gw_cls':'GW', 'ot_cls':'COPT', 'combine_cls':'[COPT + GW]'}
    plt.title('{} Predictions'.format(name2label[name]), fontsize=20)
    plt.savefig(path)
    print('fig saved to ', path)


def plot_search_acc():
    plt.clf()
    x_l = [1, 3, 5, 10, 15]
    ot_acc = [0.9721962,0.9866296,0.9941481333,0.998,0.998]
    svd_acc = [0.814787,0.894257037,0.9349997,0.9774814073,0.9888888777]
    ot_std = [0.005516745247,0.003174584773,0.0002565744596,0.003464101615,0.003464101615] 
    svd_std = [0.01152346983,0.02409426815,0.02220130725,0.009109368462,0.009622514205] 

    fig = plt.figure()
    
    
    plt.errorbar(x_l, ot_acc, yerr=ot_std, marker='o', label='COPT sketches')
    
    plt.errorbar(x_l, svd_acc, yerr=svd_std, marker='+', label='Spectral projections')
    
    plt.title('Classification acc of [COPT, GW] vs [spectral projections, GW] pipelines')
    plt.legend()
    plt.xlabel('Number of candidates allowed to 2nd stage')
    plt.ylabel('Classification accuracy')
    path = 'data/search_acc.jpg'
    fig.savefig(path)
    print('fig saved to ', path)
    

    
def create_dir(path):
    if not os.path.exists(path):
        os.mkdir(path)
    


def normalizedMI(ar1, ar2):
    score = sklearn.metrics.normalized_mutual_info_score(ar1, ar2)
    return score


def remove_edges(L, n_remove=1, seed=None):
    rng = np.random.RandomState(seed)
    
    G = lap_to_graph(L)
    edges = np.triu(L, k=1).nonzero()
    
    removed = 0
    for idx in rng.permutation(edges[0].size):
        u, v = edges[0][idx], edges[1][idx]

        G.remove_edge(u,v)
        if nx.is_connected(G):
            removed += 1
        else:
            G.add_edge(u,v)
        if removed == n_remove:
            break
     
    return graph_to_lap(G) 

def permute_nodes(l1, seed=None):
    '''
    Adapted from GOT.
    '''
    np.random.seed(seed)
    n = len(l1)
    idx = np.random.permutation(n)
    P_true = np.eye(n)
    P_true = P_true[idx]
    l2 = np.array(P_true @ l1 @ P_true.T)
    
    return np.double(l2), idx 

def symmetrize(mx, inplace=True):
    m, n = mx.size()
    assert m == n    
    mask = torch.ones(m, m)
    mask = torch.tril(mask, diagonal=-1)    
    
    if not inplace:
        mx = mx.clone()
    mx[mask > 0] = torch.triu(mx, diagonal=1).t()[mask > 0]        
    return mx

def remove_isolates(g):
    g.remove_nodes_from(list(nx.isolates(g)))

def load_data(fname):    
    with open(fname, 'rb') as f:
        data = pickle.load(f)
        graphs = data['graphs']
        labels = data['labels']
    return graphs, labels

def read_lines(path):
    with open(path, 'r') as file:
        return file.readlines()

def parse_cls(st):
    ar = st.split('., ')
    return [int(i) for i in ar]
    
def plot_confusions():
    
    ot_cls = '0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 4., 1., 1., 4., 1., 1., 4., 1., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 5., 5., 5., 5., 5., 5., 5., 5., 5., 5., 6., 6., 6., 6., 6., 6., 6., 6., 6., 6., 6., 9., 6., 9., 9., 9., 6., 6., 6., 9'
    gw_cls = '0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 6., 1., 1., 1., 1., 6., 6., 1., 6., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 5., 5., 5., 5., 5., 5., 5., 5., 5., 5., 6., 5., 6., 6., 6., 6., 5., 6., 6., 6., 9., 9., 9., 9., 9., 9., 9., 9., 9., 9'
    combine_cls = '0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 4., 1., 1., 4., 1., 1., 1., 1., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 5., 5., 5., 5., 5., 5., 5., 5., 5., 5., 6., 5., 6., 6., 6., 6., 5., 6., 6., 6., 9., 9., 9., 9., 9., 9., 9., 9., 9., 9'
    ot_cls = parse_cls(ot_cls)
    gw_cls = parse_cls(gw_cls)
    combine_cls = parse_cls(combine_cls)
    
    tgt_cls = []
    for i in [0, 1, 4, 5, 6, 9]:
        tgt_cls.extend([i]*10)
    names = ['ot_cls', 'gw_cls', 'combine_cls']
    for i, pred in enumerate([ot_cls, gw_cls, combine_cls]):
        plot_confusion(tgt_cls, pred, name=names[i])

def plot_convergence():
    '''
    plot Convergence
    '''
    conv = '99.77512221820409, 47.0570797352225, 45.454712183248795, 43.46666639511483, 40.9052209134619, 34.46097733377226, 24.431260224317242, 23.51762172579995, 22.485459983772515, 22.106446259828203, 21.534230884767375, 21.463743674621867'
    ar = conv.split(', ')
    conv_ar = [float(f) for f in ar]
    x_l = [20*i for i in list(range(len(conv_ar)))]
    fig = plt.figure()
    
    plt.plot(x_l, conv_ar, '-o', label='COPT distance')
    
    
    plt.title('COPT distance convergence sketching a 50-node graph to 15 nodes')
    plt.legend()
    plt.xlabel('Number of iterations')
    plt.ylabel('COPT distance')
    path = 'data/convergence.jpg'
    fig.savefig(path)
    print('fig saved to ', path)
    
if __name__ == '__main__':
    """
    For testing utils functions.
    """
    '''
    n = 2
    L = create_graph_lap(n)
    view_graph(L, soft_edge=True)
    '''
    
    plot_cls_acc()