Theano simple implementation of exmaple 1 in the MCMC book

theano_scan_example2.py

#!/usr/bin/python3
'''Theano simple implementation of exmaple 1 in the MCMC book, 
the same progrm with
example1.py (https://gist.github.com/GM3D/51965c9e6de5456971b5).
without shared variable.'''

from __future__ import print_function
import numpy as np
import theano
from theano.ifelse import ifelse
from theano import function, scan, shared
from theano import tensor as T

def pm_ones(N):
    '''used only for initial value of x'''
    return 2 * np.random.random_integers(0, 1, size=N) - 1

def flip(x, i):
    return T.set_subtensor(x[i], -x[i])

def f1(x):
    return x[0]

def f2(x):
    return x[0]*x[1]

def OneStep(x, i, sum_1, sum_2, theta):
    r = T.exp(-2.0*theta*x[0]*(x[1] + x[2]))
    R = trng.uniform()
    x = ifelse(T.lt(R, r), flip(x, i%N), x)
    i = i + 1
    return x, i, sum_1 + f1(x), sum_2 + f2(x)

N_value = 3
n_steps_value = 10000
initial_i = 0
initial_x = pm_ones(N_value)
theta_value = 1.0

trng = T.shared_randomstreams.RandomStreams(1234)

N = T.iscalar('N')
n_steps=T.iscalar('n_steps')
x = T.vector('x')
i = T.iscalar('i')
theta = T.fscalar('theta')
f1_sum = T.dscalar('f1_sum')
f2_sum = T.dscalar('f2_sum')

result, updates = scan(OneStep, 
                       outputs_info=[T.ones_like(x),
                                     T.zeros_like(i),
                                     T.zeros_like(f1_sum),
                                     T.zeros_like(f2_sum)],
                       non_sequences=theta,
                       n_steps=n_steps)

metropolis = function(inputs=[N, x, i, f1_sum, f2_sum, theta, n_steps],
                      outputs=(result[2][-1]/n_steps,
                               result[3][-1]/n_steps))

output = metropolis(N_value, initial_x, initial_i, 0.0, 0.0, 
                    theta_value, n_steps_value)
print('theta = %f, n_steps = %d'%(theta_value, n_steps_value))
print('result (average x1, average x1*x2) = ', map(np.ndarray.tolist, output))