changing non-defined prior and source specification

1 files changed, 6 insertions, 6 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 = {}