Parallel tempering with theano

HMC_repmon.py

# -*- coding: utf-8 -*-
import numpy
from theano import function, shared
from theano import tensor as T
import theano

import numpy as np
from scipy import linalg

sharedX = lambda X, name: shared(numpy.asarray(X, dtype=theano.config.floatX), name=name)
import exchange

def kinetic_energy(v):
    return 0.5 * (v ** 2).sum(axis=1)

def hamiltonian(x, v, logP):
    # assuming mass = 1
    return logP(x) + kinetic_energy(v)
    
def MHorg(E, En, rng):
    return (T.exp(E - En) - rng.uniform(size=E.shape)) >= 0

def MH(E, En, beta,rng):
    return (T.exp(beta*(E - En)) - rng.uniform(size=E.shape)) >= 0

def gradsum(E,x):
    return T.grad(E(x).sum(),x)
    
def dynamics(x0, v0, stepsize, n_steps, logP):
    def leapfrog(x, v, step):
        vn=v-step*gradsum(logP,x)
        xn=x+vn*step
        return [xn,vn], {}

    v_half=v0-stepsize/2*gradsum(logP,x0)
    xn=x0+stepsize*v_half

    (xs, vs), scan_updates = theano.scan(leapfrog,
            outputs_info=[
                dict(initial=xn),
                dict(initial=v_half),
                ],
            non_sequences=[stepsize],
            n_steps=n_steps - 1)

    final_x,final_v = xs[-1],vs[-1]

    assert not scan_updates
    E=logP(final_x)
    final_v = final_v - stepsize/2 * T.grad(E.sum(), final_x)

    return final_x, final_v
    
def hmc_move(rng, x0, logP, beta, stepsize, n_steps):
    v0 = rng.normal(size=x0.shape)
    
    xn,vn = dynamics(
            x0=x0,  v0=v0,
            stepsize=stepsize,  n_steps=n_steps,
            logP=logP)

    E =hamiltonian(x0, v0, logP)
    En=hamiltonian(xn, vn, logP)
    accept = MH(E,En,beta,rng=rng)
    
    return accept,xn,T.switch(accept,En,E)
        
def hmc_updates(x, stepsize, avg_acceptance_rate, final_x, accept,
                 target_acceptance_rate, stepsize_inc, stepsize_dec,
                 stepsize_min, stepsize_max, avg_acceptance_slowness):

    ## POSITION UPDATES ##
    accept_matrix = accept.dimshuffle(0, *(('x',) * (final_x.ndim - 1)))
    xn = T.switch(accept_matrix, final_x, x)

    ## STEPSIZE UPDATES ##
    _new_stepsize = T.switch(avg_acceptance_rate > target_acceptance_rate,
                              stepsize * stepsize_inc, stepsize * stepsize_dec)
    new_stepsize  = T.clip(_new_stepsize, stepsize_min, stepsize_max)

    ## ACCEPT RATE U+PDATES ##
    mean_dtype = theano.scalar.upcast(accept.dtype, avg_acceptance_rate.dtype)
    new_acceptance_rate = T.add(
            avg_acceptance_slowness * avg_acceptance_rate,
            (1.0 - avg_acceptance_slowness) * accept.mean(dtype=mean_dtype))

    return [(x, xn),
            (stepsize, new_stepsize),
            (avg_acceptance_rate, new_acceptance_rate)]
    

def onerun(xs_shared,rng,
                logP,
                beta,
                stepsize,
                n_steps,
                avg_acceptance_rate,
                step_min,
                step_max,
                step_inc,
                step_dec,
                target_acceptance_rate,
                avg_acceptance_slowness):

       accept, final_pos, finalE= hmc_move(
                rng,
                xs_shared,
                logP,
                beta,
                stepsize,
                n_steps)
    
       xs,stepsizes,acceptrates = hmc_updates(
                xs_shared,
                stepsize,
                avg_acceptance_rate,
                final_x=final_pos,
                accept=accept,
                stepsize_min=step_min,
                stepsize_max=step_max,
                stepsize_inc=step_inc,
                stepsize_dec=step_dec,
                target_acceptance_rate=target_acceptance_rate,
                avg_acceptance_slowness=avg_acceptance_slowness)
 
       return finalE,xs,stepsizes,acceptrates

def onerunexchange(xs_shared,Es_shared,rng,
                logP, beta, indices, parity,
                stepsize,  n_steps,
                avg_acceptance_rate,
                step_min, step_max,
                step_inc, step_dec,
                target_acceptance_rate,
                avg_acceptance_slowness):
            
            Es,xs,stepsizes,acceptrates=onerun(
                    xs_shared,rng,
                    logP,beta,
                    stepsize,
                    n_steps,
                    avg_acceptance_rate,
                    step_min,
                    step_max,
                    step_inc,
                    step_dec,
                    target_acceptance_rate,
                    avg_acceptance_slowness
                    )
                    
            xs,u=exchange.exchange(xs,Es,beta,indices,parity)
            xs_update = (xs_shared, T.set_subtensor(xs_shared[::], xs))

            return [(xs_shared,xs_update),Es,stepsizes,acceptrates]
            
def onerunexchange(xs_shared,Es_shared,rng,
                logP, beta, indices, parity,
                stepsize,  n_steps,
                avg_acceptance_rate,
                step_min, step_max,
                step_inc, step_dec,
                target_acceptance_rate,
                avg_acceptance_slowness):

            xs,u=exchange.exchange(xs_shared,Es_shared,beta,indices,parity)
            xs_update = (xs_shared, T.set_subtensor(xs_shared[::], xs))
            return [(xs_shared,xs_update),(step_min,step_max)]
    
class HMC_sampler(object):

    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)

    @classmethod
    def new(cls, xs_shared,Es,logP,betas,batchsize,parity,
            initial_stepsize=0.01, target_acceptance_rate=.9, n_steps=20,
            step_dec=0.98,
            step_min=0.001,
            step_max=0.25,
            step_inc=1.02,
            avg_acceptance_slowness=0.9,
            seed=12345):

        betas=T.dvector("beta")                

       #[n_samples, batchsize, dim]
 #       batchsize = xs_shared.shape[0]
        indices=shared(np.array(range(0,batchsize,2)))
        stepsize = sharedX(initial_stepsize, 'hmc_stepsize')
        avg_acceptance_rate = sharedX(target_acceptance_rate, 'avg_acceptance_rate')

        rng = T.shared_randomstreams.RandomStreams(seed)
        updates=onerunexchange(xs_shared,Es,rng,
                                  logP, betas, indices, parity,
                stepsize,  n_steps,
                avg_acceptance_rate,
                step_min, step_max,
                step_inc, step_dec,
                target_acceptance_rate,
                avg_acceptance_slowness)

        # compile
        simulate = function([], [], updates=updates)

        return cls(
                positions=xs_shared,
                Es=Es,
                stepsize=stepsize,
                stepsize_min=step_min,
                stepsize_max=step_max,
                avg_acceptance_rate=avg_acceptance_rate,
                target_acceptance_rate=target_acceptance_rate,
                beta=betas,
                rng=rng,
                _updates=updates,
                simulate=simulate)

#1 set with exchange
    def draw(self, **kwargs):
        self.simulate()
        return self.positions.get_value(borrow=False),self.Es.get_value(borrow=False)



dim=3
seed=120

rng = np.random.RandomState(seed)
mu  = np.array(rng.rand(dim) * 10, dtype=theano.config.floatX)
cov = np.array(rng.rand(dim, dim), dtype=theano.config.floatX)
cov = (cov + cov.T) / 2.
cov[np.arange(dim), np.arange(dim)] = 1.0
cov_inv = linalg.inv(cov)

gaussian_energy=lambda x:(T.dot((x - mu), cov_inv) * (x - mu)).sum(axis=1)/2
_gaussian=lambda x,mu,cov_inv:(T.dot((x - mu), cov_inv) * (x - mu)).sum(axis=1)/2

mixgaussianfunc=lambda x,mu,covinv: _gaussian(mu[0],covinv[0])+_gaussian(mu[1],covinv[1])

def sampler_on_nd_gaussian(sampler_cls, burnin, n_samples,dim=10):
    batchsize=10

    rng = np.random.RandomState(123)
    # Declared shared random variable for positions
    x = shared(rng.randn(batchsize, dim).astype(theano.config.floatX))
    Es=shared(rng.randn(dim).astype(theano.config.floatX))
    betas=shared(np.array(range(batchsize)))
#    indices=shared(range(0,batchsize,2))
    parity=shared(1)
    
    # Create HMC sampler
    sampler = sampler_cls(x, Es,gaussian_energy,betas,batchsize,parity,
            initial_stepsize=1e-3, step_max=0.5)

    [sampler.draw() for r in xrange(burnin)]  #burn-in
    # Draw `n_samples`: result is a 3D tensor of dim [n_samples, batchsize, dim]
    _samples = np.asarray([sampler.draw() for r in xrange(n_samples)])
    # Flatten to [n_samples * batchsize, dim]
    samples = _samples.T.reshape(dim, -1).T

    print 'target mean:', mu
    print 'target cov:\n', cov

    print 'empirical mean: ', samples.mean(axis=0)
    print 'empirical_cov:\n', np.cov(samples.T)

    print 'final stepsize', sampler.stepsize.get_value()
    print 'final acceptance_rate', sampler.avg_acceptance_rate.get_value()

    return sampler

def test_hmc():
    sampler = sampler_on_nd_gaussian(HMC_sampler.new,
            burnin=1000, n_samples=1000, dim=5)
    assert abs(sampler.avg_acceptance_rate.get_value() -
               sampler.target_acceptance_rate) < .1
    assert sampler.stepsize.get_value() >= sampler.stepsize_min
    assert sampler.stepsize.get_value() <= sampler.stepsize_max
    
if __name__ == "__main__":
    test_hmc()