diff options
Diffstat (limited to 'simulation')
| -rw-r--r-- | simulation/bayes.py | 12 | ||||
| -rw-r--r-- | simulation/main.py | 6 |
2 files changed, 9 insertions, 9 deletions
diff --git a/simulation/bayes.py b/simulation/bayes.py index 7d8567b..e2be6da 100644 --- a/simulation/bayes.py +++ b/simulation/bayes.py @@ -45,13 +45,13 @@ def mc_graph_setup(infected, susceptible, prior=None, *args, **kwargs): # Container class for graph parameters n_nodes = len(infected[0][0]) theta = np.empty((n_nodes,n_nodes), dtype=object) - for i in xrange(n_nodes): - for j in xrange(n_nodes): - if prior is None: + if prior is None: + for i in xrange(n_nodes): + for j in xrange(n_nodes): theta[i, j] = pymc.Beta('theta_{}{}'.format(i, j), alpha=1, - beta=1) - else: - theta[i, j] = prior('theta_{}{}'.format(i, j), *args, **kwargs) + beta=1) + else: + theta = prior(theta=theta, *args, **kwargs) # Container class for cascade realization x = {} diff --git a/simulation/main.py b/simulation/main.py index e4b30f2..3e3de4b 100644 --- a/simulation/main.py +++ b/simulation/main.py @@ -108,15 +108,15 @@ def simulate_cascade(x, graph): yield x, susc -def uniform_source(graph): +def uniform_source(graph, *args, **kwargs): x0 = np.zeros(graph.shape[0], dtype=bool) x0[nr.randint(0, graph.shape[0])] = True return x0 def simulate_cascades(n, graph, source=uniform_source): - for _ in xrange(n): - x0 = source(graph) + for t in xrange(n): + x0 = source(graph, t) yield simulate_cascade(x0, graph) |