From dae606656be0302c97902bf982ad9abaf5a34d1c Mon Sep 17 00:00:00 2001
From: jeanpouget-abadie <jean.pougetabadie@gmail.com>
Date: Tue, 10 Nov 2015 17:29:10 -0500
Subject: adding formating issue when n_nodes > 10

---
 simulation/bayes.py | 25 ++++++++++++++++++-------
 1 file changed, 18 insertions(+), 7 deletions(-)

(limited to 'simulation/bayes.py')

diff --git a/simulation/bayes.py b/simulation/bayes.py
index e2be6da..82dc0c1 100644
--- a/simulation/bayes.py
+++ b/simulation/bayes.py
@@ -35,6 +35,10 @@ def glm_node_setup(cascade, y_obs, prior=None, *args, **kwargs):
     return pymc.Model([y, pymc.Container(theta), pymc.Container(x)])
 
 
+def formatLabel(s, n):
+    return '0'*(len(str(n)) - len(str(s))) + str(s)
+
+
 def mc_graph_setup(infected, susceptible, prior=None, *args, **kwargs):
     """
     Build an IC PyMC graph-level model from:
@@ -42,14 +46,16 @@ def mc_graph_setup(infected, susceptible, prior=None, *args, **kwargs):
         -susceptible nodes over time: same format as above
     Note: we use the Markov Chain formulation: X_{t+1}|X_t,theta = f(X_t.theta)
     """
+
     # Container class for graph parameters
     n_nodes = len(infected[0][0])
     theta = np.empty((n_nodes,n_nodes), dtype=object)
     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)
+                theta[i, j] = pymc.Beta('theta_{}{}'.format(formatLabel(i,
+                                        n_nodes), formatLabel(j, n_nodes)),
+                                        alpha=1, beta=1)
     else:
         theta = prior(theta=theta, *args, **kwargs)
 
@@ -78,22 +84,27 @@ if __name__=="__main__":
     g = np.log(1. / (1 - p * g))
 
     print('running the graph-level MC set-up')
+    n_nodes = len(g)
     cascades = main.simulate_cascades(1000, g)
     infected, susc = main.build_cascade_list(cascades)
     model = mc_graph_setup(infected, susc)
     mcmc = pymc.MCMC(model)
     mcmc.sample(10**4, 1000)
     fig, ax = plt.subplots(len(g), len(g))
-    for i in xrange(len(g)):
-        for j in xrange(len(g)):
-            ax[i,j].locator_params(nbins=3, axis='x')
-            ax[i,j].hist(mcmc.trace('theta_{}{}'.format(i,j))[:], normed=True)
+    for i in xrange(n_nodes):
+        for j in xrange(n_nodes):
+            if n_nodes < 5:
+                ax[i,j].locator_params(nbins=3, axis='x')
+            else:
+                ax[i, j].get_xaxis().set_ticks([])
+                ax[i, j].get_yaxis().set_ticks([])
+            it, jt = formatLabel(i, n_nodes), formatLabel(j, n_nodes)
+            ax[i,j].hist(mcmc.trace('theta_{}{}'.format(it,jt))[:], normed=True)
             ax[i, j].set_xlim([0,1])
             ax[i, j].plot([g[i, j]]*2, [0, ax[i,j].get_ylim()[-1]], color='red')
     plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=.1)
     plt.show()
 
-
     print('running the node level set-up')
     node = 0
     cascades = main.simulate_cascades(100, g)
-- 
cgit v1.2.3-70-g09d2