changing non-defined prior and source specification

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)