cascade-reconstruction-by-tree-samples/experiment_obs_method.py at master · xiaohan2012/cascade-reconstruction-by-tree-samples · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
# coding: utf-8

import os
import pandas as pd

from graph_tool import load_graph
from glob import glob
from tqdm import tqdm
from joblib import Parallel, delayed
from itertools import product

from eval_helpers import eval_map
from experiment import one_run

infection_proba = 0.1

methods = ['our', 'pagerank', 'min-steiner-tree']
# methods = ['our']
root_sampler_names = ['true_root']
obs_method = 'bfs-head'
# root_sampler_names = [None]

cascade_model = 'si'

# a batch of settings to iterate through
settings = [
    # for grqc
    # {'graphs': ['grqc'],
    #  'obs_fractions': ["0.5"],
    #  'cascade_fractions': ["0.05", "0.1", "0.15", "0.2", "0.25"]},
    # {'graphs': ['grqc'],
    #  'obs_fractions': ["0.5", "0.6", "0.7", "0.8", "0.9"],
    #  'cascade_fractions': ["0.05"]},

    # # for lattice and infectious
    # {'graphs': ['infectious', 'lattice-1024'],
    #  'obs_fractions': ["0.5"],
    #  'cascade_fractions': ["0.1", "0.2", "0.3", "0.4", "0.5"]},
    {'graphs': ['infectious', 'lattice-1024', 'grqc'],
     'obs_fractions': ["0.1", "0.2", "0.3", "0.4", "0.5", "0.6", "0.7", "0.8", "0.9"],
     'cascade_fractions': ["0.1"]}
]

for setting in settings:
    graphs, obs_fractions, cascade_fractions = setting['graphs'], \
                                               setting['obs_fractions'], \
                                               setting['cascade_fractions']
    for graph, obs_fraction, cascade_fraction, method, root_sampler_name \
            in product(
                graphs, obs_fractions, cascade_fractions, methods, root_sampler_names
            ):
        g = load_graph('data/{}/graph_weighted_{}.gt'.format(graph, infection_proba))
        edge_weights = g.edge_properties['weights']

        dataset_id = "{}-m{}-s{}-o{}-om{}".format(graph, cascade_model, cascade_fraction, obs_fraction, obs_method)
        print('method', method)
        print('dataset_id', dataset_id)

        input_dir = 'cascade/{}/'.format(dataset_id)

        if root_sampler_name is not None and method == 'our':
            output_dir = 'output/{}-{}/{}/'.format(method, root_sampler_name, dataset_id)
            eval_result_path = 'eval/{}-{}/{}.pkl'.format(method, root_sampler_name, dataset_id)
        else:
            output_dir = 'output/{}/{}/'.format(method, dataset_id)
            eval_result_path = 'eval/{}/{}.pkl'.format(method, dataset_id)

        if not os.path.exists(os.path.dirname(eval_result_path)):
            os.makedirs(os.path.dirname(eval_result_path))

        rows = Parallel(n_jobs=-1)(delayed(one_run)(g, edge_weights, input_path, output_dir, method,
                                                    root_sampler_name=root_sampler_name)
                                   for input_path in tqdm(glob(input_dir + '*.pkl')))
        assert len(rows) > 0, 'nothing calculated'

        scores = eval_map(input_dir, output_dir)

        summ = pd.Series(scores).describe()
        print(summ)
        summ.to_pickle(eval_result_path)