link_prediction_eval.py

# coding = utf-8

from evalne.evaluation.evaluator import LPEvaluator
from evalne.evaluation.split import EvalSplit as LPEvalSplit
from evalne.utils import preprocess as pp

from .lambda_func import hash_func

def get_auc_heuristics(G,timeout=60):
    H, _ = pp.prep_graph(G.copy(),maincc=True)
    traintest_split = LPEvalSplit()
    traintest_split.compute_splits(H, split_alg="spanning_tree", train_frac=0.90, fe_ratio=1)
    nee = LPEvaluator(traintest_split)
    auc_spatial, auc_sbm = 0, 0
    try:
        auc_spatial = nee.evaluate_baseline(method="spatial_link_prediction",timeout=timeout).test_scores.auroc()
        auc_sbm = nee.evaluate_baseline(method="stochastic_block_model",timeout=timeout).test_scores.auroc()
    except Exception as e:
        print("Could not compute AUC ! ")
        print(e)
    return auc_sbm,auc_spatial


def split_train_test(G,train_frac=0.90,fe_ratio=1):
    H, _ = pp.prep_graph(G.copy(), maincc=True,relabel=False)
    traintest_split = LPEvalSplit()
    traintest_split.compute_splits(H, split_alg="spanning_tree", train_frac=train_frac, fe_ratio=fe_ratio)

    X_train = traintest_split.train_edges
    y_train = traintest_split.train_labels
    X_test = traintest_split.test_edges
    y_test = traintest_split.test_labels

    return X_train,X_test,y_train,y_test

def get_all_possible_edges(G):
    register = set([])
    data = []
    for n1 in list(G.nodes()):
        for n2 in list(G.nodes()):
            if n1 != n2 and hash_func((n1, n2)) not in register:
                data.append([n1, n2])
                register.add(hash_func((n1, n2)))
    return data