evalNE_script.py

from evalne.evaluation.evaluator import LPEvaluator
from evalne.evaluation.split import EvalSplit as LPEvalSplit
from evalne.evaluation.score import Scoresheet
from evalne.utils import preprocess as pp
import networkx as nx

from lib.utils import load_edgelist

import argparse

parser = argparse.ArgumentParser()
parser.add_argument("edgelist_graph_filename")
parser.add_argument("--ne","--network-embedding",action="store_true",help="If you want to use neural network embedding for link prediction")
parser.add_argument("-v","--verbose",action="store_true")

args = parser.parse_args()#("data/fb_country_country_sample_6_size1000.txt".split())

# Load and preprocess the network
G = nx.read_gexf(args.edgelist_graph_filename)#load_edgelist(args.edgelist_graph_filename,is_directed=True,weighted=True)
G, _ = pp.prep_graph(G,maincc=True)


# Create an evaluator and generate train/test edge split
traintest_split = LPEvalSplit()
traintest_split.compute_splits(G,split_alg="spanning_tree",train_frac=0.8,fe_ratio=1)
nee = LPEvaluator(traintest_split)
# Create a Scoresheet to store the results
scoresheet = Scoresheet()

# Set the baselines
methods = ['random_prediction',
    'common_neighbours',
    'jaccard_coefficient',
    "adamic_adar_index",
    "preferential_attachment",
    "resource_allocation_index",
    "stochastic_block_model",
    "stochastic_block_model_edge_probs",
    "stochastic_block_model_degree_corrected"
           ]

# Evaluate baselines
for method in methods:
    result = nee.evaluate_baseline(method=method, )
    scoresheet.log_results(result)

if args.network_embedding:
    try:
        # Check if OpenNE is installed
        import openne
        # Set embedding methods from OpenNE
        methods = "node2vec hope-opne gf sdne deepWalk line grarep".split() #lap-opne
        commands = [
            "python -m openne --method node2vec --graph-format edgelist --epochs 100 --number-walks 10 --walk-length 80 --window-size 10",
            "python -m openne --method hope --epochs 100",
            "python -m openne --method gf --epochs 100",
            "python -m openne --method sdne --epochs 100 --encoder-list [1024,128] --beta 5 --bs 500",
            "python -m openne --method deepWalk --graph-format edgelist --epochs 100 --number-walks 10 --walk-length 80 --window-size 10",
            "python -m openne --method line --graph-format edgelist --epochs 10",
            "python -m openne --method grarep --epochs 100"
            # "python -m openne --method lap --epochs 100",
            ]
        edge_emb = ['average', 'hadamard']

        # Evaluate embedding methods
        for i in range(len(methods)):
            command = commands[i] + " --input {} --output {} --representation-size {}"
            results = nee.evaluate_cmd(method_name=methods[i], method_type='ne', command=command,
                                       edge_embedding_methods=edge_emb, input_delim=' ', output_delim=' ',  verbose=args.verbose)
            scoresheet.log_results(results)

    except ImportError:
        print("The OpenNE library is not installed. Reporting results only for the baselines...")
        pass

# Get output
if args.verbose:
    scoresheet.print_tabular()
scoresheet.write_all(args.edgelist_graph_filename+"_results_lp")