diff options
| -rw-r--r-- | simulation/bayes.py | 25 |
1 files changed, 18 insertions, 7 deletions
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) |