dTRAM_example.ipynb

{
 "metadata": {
  "name": "",
  "signature": "sha256:32cf7fd884672fd482711c6d5701c3ac6bb65f00402928831be84bea85cdb445"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "heading",
     "level": 1,
     "metadata": {},
     "source": [
      "How to use the dTRAM estimator"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "In this short tutorial, we will learn use the dTRAM method to estimate stationary distributions and transition matrices from multi-ensemble simulations.\n",
      "\n",
      "We start by importing the TRAMData class and dTRAM's API function from the pytram package."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%pylab inline\n",
      "from pytram import TRAMData, dtram # this is the dTRAM API function"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "heading",
     "level": 2,
     "metadata": {},
     "source": [
      "First example: three discrete states with piecewise constant potentials"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "The first example uses a system with three discrete states (L,C,R), where each state corresponds to a certain energy. We allow transitions between states L and C and between C and R. Furthermore, the state C corresponds to a higher potential energy than states L and R. Thus, state C acts as a transition state between the metastable states L and R.\n",
      "\n",
      "We perform two unbiased Metropolis Monte Carlo simulations starting in each metastable state and an additional biased simulation starting from the transition state. The bias raises the energies of the metastable states and allows for free diffusion between L and C, and C and R."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# piecewise contant potentials\n",
      "u_i = np.array( [ 4.0, 8.0, 0.0 ], dtype=np.float64 )\n",
      "# bias energies\n",
      "b_K_i = np.array( [ [ 0.0, 0.0, 0.0 ], [ 4.0, 0.0, 8.0 ] ], dtype=np.float64 )\n",
      "# proposal probabilities\n",
      "S = np.ones( shape=(3,3), dtype=np.float64 ) * 0.5\n",
      "S[0,2] = 0.0\n",
      "S[1,1] = 0.0\n",
      "S[2,0] = 0.0\n",
      "# transition matrix generator\n",
      "def tmat( s, e ):\n",
      "    a = np.zeros( shape=s.shape, dtype=np.float64 )\n",
      "    for i in xrange( e.shape[0] ):\n",
      "        for j in xrange( e.shape[0] ):\n",
      "            a[i,j] = e[j] - e[i]\n",
      "    idx = ( a < 0.0 )\n",
      "    a[idx] = 0.0\n",
      "    a = np.exp( -a )\n",
      "    T = s*a\n",
      "    for i in xrange( T.shape[0] ):\n",
      "        T[i,i] = 0.0\n",
      "        T[i,i] = 1.0 - T[i,:].sum()\n",
      "    return T\n",
      "# unbiased transition matrix at kT=1.0\n",
      "Tu = tmat( S, u_i+b_K_i[0,:] )\n",
      "# biased transition matrix at kT=1.0\n",
      "Tb = tmat( S, u_i+b_K_i[1,:] )\n",
      "# throw an index dice according to the distribution T_i\n",
      "def n_dice( T_i ):\n",
      "    rnd = np.random.rand()\n",
      "    cdf = np.cumsum( T_i )\n",
      "    ind = ( cdf > rnd )\n",
      "    idx = np.where( ind == True )\n",
      "    return np.min( idx )\n",
      "# get a markov chain according to a transition matrix T, staring from state i with length L\n",
      "def m_chain( T, i, L ):\n",
      "    dtraj = np.zeros( shape=(L,), dtype=np.intc )\n",
      "    dtraj[0] = i\n",
      "    for l in xrange( 1, L ):\n",
      "        dtraj[l] = n_dice( T[dtraj[l-1],:] )\n",
      "    return dtraj\n",
      "# get the sequential data from a single simulation\n",
      "def get_sequence( T, i, L, therm_state ):\n",
      "    m = m_chain( T, i, L )\n",
      "    t = np.ones( shape=m.shape, dtype=np.intc ) * therm_state\n",
      "    return { 'm': m, 't': t }\n",
      "# run the simulations of length L\n",
      "L = 10000\n",
      "T0 = get_sequence( Tu, 0, L, 0 )\n",
      "T1 = get_sequence( Tu, 2, L, 0 )\n",
      "T2 = get_sequence( Tb, 1, L, 1 )"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "In the next step, we prepare our input data for dTRAM, i.e., we let the TRAMData object handle the preparation."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "tramdata = TRAMData( [T0,T1,T2], b_K_i=b_K_i )"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Now, we call the dTRAM API function which estimates the unbiased staionary probabilities of the three states L, C, and R."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "dtram_obj = dtram( tramdata, 1, maxiter=100000, ftol=1.0E-15, verbose=False )"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "If we have not made any mistakes, the dtram_obj should now carry the stationary distribution, which we will compare against the exact solution."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "pi_exact = np.exp( -u_i )\n",
      "pi_exact /= pi_exact.sum()\n",
      "print \"pi_estimated               :\", dtram_obj.pi_i\n",
      "print \"pi_exact                   :\", pi_exact\n",
      "print \"pi_estimated normalization :\", dtram_obj.pi_i.sum()\n",
      "print \"||pi_estimated - pi_exact||:\", np.linalg.norm( dtram_obj.pi_i - pi_exact )"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Furthermore, we can also estimate a transition matrix for every thermodynamic state."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "T_estimated = dtram_obj.estimate_transition_matrices()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Finally, we compare the estimated transition matrices against the exact ones."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "print \"T_estimated_unbiased - Tu:\"\n",
      "print (T_estimated[0,:,:] - Tu)\n",
      "print \"||T_estimated_unbiased - Tu||:\", np.linalg.norm( T_estimated[0,:,:] - Tu )\n",
      "print \"T_estimated_biased - Tb:\"\n",
      "print (T_estimated[1,:,:] - Tb)\n",
      "print \"||T_estimated_biased - Tb||:\", np.linalg.norm( T_estimated[1,:,:] - Tb )"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "heading",
     "level": 2,
     "metadata": {},
     "source": [
      "Second example: symmetric double well potential"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "The second example addresses a double well potential in one dimension, i.e., we have still two metastable states which are separated by a potential barrier, but we do not have piecewise constant potentials anymore.\n",
      "\n",
      "Again, we use the Metropolis Monte Carlo approach to generate trajectories, which have additional harmonic restraints centered at several points across the potential barrier to increase the number of transition events. This is called the umbrella sampling approach."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# harmonic bias potential\n",
      "def B( x, xk ):\n",
      "    if None == xk:\n",
      "        return 0.0\n",
      "    return 4.0*( x - xk )**2\n",
      "# double well potential (with otiopnal bias)\n",
      "def U( x, xk=None ):\n",
      "    return 0.25*x**4 - 5.0*x**2 - 9.9874 + B( x, xk )\n",
      "# Metropolis Monte Carlo\n",
      "def mmc( x0, L, xk=None ):\n",
      "    T = [x0]\n",
      "    E = U( x0, xk )\n",
      "    for t in xrange( L ):\n",
      "        x = T[-1] + 0.4*( np.random.rand() - 0.5 )\n",
      "        e = U( x, xk )\n",
      "        accept = False\n",
      "        if e < E:\n",
      "            accept = True\n",
      "        elif np.random.rand() < np.exp( -( e - E ) ):\n",
      "            accept = True\n",
      "        if accept:\n",
      "            T.append( x )\n",
      "            E = e\n",
      "        else:\n",
      "            T.append( T[-1] )\n",
      "    return np.array( T, dtype=np.float64 )\n",
      "# discretisation scheme\n",
      "def discretise( T, centers ):\n",
      "    D = np.zeros( shape=T.shape, dtype=np.int32 )\n",
      "    for i in xrange( T.shape[0] ):\n",
      "        if T[i] <= centers[0]:\n",
      "            D[i] = 0\n",
      "        elif T[i] >= centers[-1]:\n",
      "            D[i] = centers.shape[0]-1\n",
      "        else:\n",
      "            for j in xrange( centers.shape[0] ):\n",
      "                if T[i] < centers[j]+0.05:\n",
      "                    D[i] = j\n",
      "                    break\n",
      "    return D\n",
      "# generate trajectories\n",
      "def get_trajectories( centers, xk, L ):\n",
      "    r = []\n",
      "    for i in xrange( xk.shape[0] ):\n",
      "        traj = mmc( xk[i], L, xk=xk[i] )\n",
      "        dtraj = discretise( traj, centers )\n",
      "        r.append( { 'm': dtraj, 't': np.ones( shape=dtraj.shape, dtype=np.intc ) * i } )\n",
      "    N_K_i = np.zeros( shape=(xk.shape[0],centers.shape[0]), dtype=np.intc )\n",
      "    for K in xrange( xk.shape[0] ):\n",
      "        for i in xrange( centers.shape[0] ):\n",
      "            N_K_i[K,i] = ( r[K]['m'] == i ).sum()\n",
      "    N_i = N_K_i.sum( axis=0 )\n",
      "    idx = ( N_i > 0 )\n",
      "    NSTATES = idx.sum()\n",
      "    for i in reversed( xrange( centers.shape[0] ) ):\n",
      "        if idx[i]:\n",
      "            continue\n",
      "        for K in xrange( xk.shape[0] ):\n",
      "            idx2 = ( r[K]['m'] > i )\n",
      "            r[K]['m'][idx2] -= 1\n",
      "    return r , idx\n",
      "# set the discretisation and simulation parameters\n",
      "NBINS = 101\n",
      "centers = np.linspace( -5.0, 5.0, NBINS )\n",
      "NBIAS = 11\n",
      "xk = np.linspace( -5.0, 5.0, NBIAS )\n",
      "L = 10000\n",
      "# run the simulations\n",
      "trajs , idx = get_trajectories( centers, xk, L )\n",
      "# calculate the number of visisted states\n",
      "NSTATES = idx.sum()\n",
      "# calculate the bias energies for the visited states\n",
      "b_K_i = np.zeros( shape=(xk.shape[0],NSTATES), dtype=np.float64 )\n",
      "for K in xrange( xk.shape[0] ):\n",
      "    b_K_i[K,:] = B( centers[idx], xk[K] )"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "In the next step, we prepare our input data for dTRAM, i.e., we let the TRAMData object handle the preparation."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "tramdata = TRAMData( trajs, b_K_i=b_K_i )"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "We call the API function..."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "dtram_obj = dtram( tramdata, 1, maxiter=100000, ftol=1.0E-4, verbose=False )"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "...and compare our estimation of the unbiased stationary probabilities via the corresponding free energies against the exactly known double well potential."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "nz = dtram_obj.pi_i.nonzero()\n",
      "cnz = (centers[idx])[nz]\n",
      "u1 = U( cnz )\n",
      "u2 = dtram_obj.f_i[nz]\n",
      "u2 -= u2.min() - u1.min()\n",
      "plt.plot( cnz, u2, '--o', color='green', label=\"dTRAM\" )\n",
      "plt.plot( cnz, u1, label=\"exact\" )\n",
      "plt.legend( loc=1, fontsize=10 )\n",
      "plt.xlabel( r\"$x$ / a.u.\", fontsize=15 )\n",
      "plt.ylabel( r\"$U(x)$ / kT\", fontsize=15 )"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "heading",
     "level": 2,
     "metadata": {},
     "source": [
      "Concluding remark"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "We hope that these examples help to understand the basic usage of the dTRAM method in the pytram package. Feel free to contact us via\n",
      "\n",
      "    pytram@lists.fu-berlin.de\n",
      "\n",
      "if any problems or questions arise."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "print \"If you find dTRAM useful, please cite\\n\"\n",
      "dtram_obj.cite()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}