def test_main(): g = lattice((5, 5)) for n, c in [(0, 1), (10, 5), (g.num_vertices(), 1)]: for _ in range(c): obs = np.random.choice(np.arange(g.num_vertices()), g.num_vertices(), replace=False) t = min_steiner_tree(g, obs) assert is_tree(t) assert set(obs).issubset(set(map(int, t.vertices())))
def test_unweighted(): g = lattice((5, 5)) for n, c in [(10, 5), (g.num_vertices(), 1)]: # repeat `c` rounds, using `n` terminals for _ in range(c): obs = np.random.choice(np.arange(g.num_vertices()), n, replace=False) t = min_steiner_tree(g, obs) assert is_tree(t) assert set(obs).issubset(set(map(int, t.vertices())))
def one_run(g, edge_weights, input_path, output_dir, method='our', **kwargs): basename = os.path.basename(input_path) output_path = os.path.join(output_dir, basename) if os.path.exists(output_path): # print(output_path, 'procssed, skip') return obs, c = pkl.load(open(input_path, 'rb')) nlog_edge_weights = g.new_edge_property('float') nlog_edge_weights.a = -np.log(edge_weights.a) if method == 'our': root_sampler_name = kwargs.get('root_sampler_name') root_sampler = get_root_sampler_by_name(root_sampler_name, g=g, obs=obs, c=c, weights=nlog_edge_weights) n_samples = kwargs.get('n_sample', 5000) inf_probas = infection_probability_shortcut(g, edge_weights=edge_weights, obs=obs, root_sampler=root_sampler, n_samples=n_samples, log=False) elif method == 'min-steiner-tree': from minimum_steiner_tree import min_steiner_tree # we want the product of weights, so apply negative log # nlog_edge_weights = g.new_edge_property('float') # nlog_edge_weights.a = -np.log(edge_weights.a) nodes = min_steiner_tree(g, obs, p=nlog_edge_weights, return_type='nodes') # make it a binary vector inf_probas = np.zeros(len(c)) inf_probas[nodes] = 1 elif method == 'pagerank': # inf_probas is not actually probability inf_probas = pagerank_scores(g, obs, weight=edge_weights) else: raise ValueError('unsupported method') pkl.dump({'inf_probas': inf_probas}, open(output_path, 'wb'))
def test_weighted(): g = lattice((5, 5)) # assign some random weight p = g.new_edge_property('float') weights = np.random.random(g.num_edges()) p.a = (-np.log(weights)) for n, c in [(10, 5), (g.num_vertices(), 1)]: # repeat `c` rounds, using `n` terminals for _ in range(c): obs = np.random.choice(np.arange(g.num_vertices()), n, replace=False) t = min_steiner_tree(g, obs, p) # print(t) assert is_tree(t) assert set(obs).issubset(set(map(int, t.vertices())))
def one_run_for_edge(g, edge_weights, input_path, output_dir, method='our', **kwargs): basename = os.path.basename(input_path) output_path = os.path.join(output_dir, basename) if os.path.exists(output_path): # print(output_path, 'procssed, skip') return obs, c, _ = pkl.load(open(input_path, 'rb')) nlog_edge_weights = g.new_edge_property('float') nlog_edge_weights.a = -np.log(edge_weights.a) if method == 'our': root_sampler_name = kwargs.get('root_sampler_name') root_sampler = get_root_sampler_by_name(root_sampler_name, g=g, obs=obs, c=c, weights=nlog_edge_weights) n_samples = kwargs.get('n_sample', 1000) edge_freq = infer_edge_frequency(g, edge_weights=edge_weights, obs=obs, root_sampler=root_sampler, n_samples=n_samples, log=False) elif method == 'min-steiner-tree': from minimum_steiner_tree import min_steiner_tree edges = min_steiner_tree(g, obs, p=nlog_edge_weights, return_type='edges') # make it a binary vector edge_freq = {e: 1 for e in edges} else: raise ValueError('unsupported method') pkl.dump({'edge_freq': edge_freq}, open(output_path, 'wb'))
def infer_infected_nodes(g, obs, estimator=None, use_proba=True, method="min_steiner_tree", min_inf_proba=0.5): """besides observed infections, infer other infected nodes if method is 'sampling', refer to infection_probability, `min_inf_proba` is the minimum infection probability to be considered "'infected' """ assert method in {"min_steiner_tree", "sampling"} if method == 'min_steiner_tree': st = min_steiner_tree(g, obs) remain_infs = set(map(int, st.vertices())) return remain_infs else: # sampling assert estimator is not None, 'sampling approach requires an estimator' inf_probas = estimator.unconditional_proba() if use_proba: return inf_probas else: return (inf_probas >= min_inf_proba).nonzero()[0]