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]