xTRAM_example.ipynb

{
 "metadata": {
  "name": "",
  "signature": "sha256:979341b2a17200789b17ee3d676dc8858ea7b417d41622926f7082c45e447d5e"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "heading",
     "level": 1,
     "metadata": {},
     "source": [
      "How to use the xTRAM estimator"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "First we will import all we need from the pytram package"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%pylab inline\n",
      "from pytram import TRAMData, xtram\n",
      "import matplotlib.pylab as plt"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "heading",
     "level": 2,
     "metadata": {},
     "source": [
      "Simulated tempering example: asymmetric double well potential"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Below we construct a Brownian dynamics simulation of an asymmetric double well potential similar to the one used in the xTRAM paper [1]. We will generate a single simulated tempering trajectory in a form that pytram will understand. All energies will be reduced, i.e. u/kT. "
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# assymetric double well potential \n",
      "def U( x ):\n",
      "    return 2*(x-2) - 6*(x-2)**2 + (x-2)**4\n",
      "# gradient of the potential\n",
      "def G( x ):\n",
      "    return 4*x**3-24*x**2+36*x-6\n",
      "# simulated tempering temperature exchange step\n",
      "def change_temp( ct, E, temp_array,f):\n",
      "    newt = np.random.randint(temp_array.shape[0])\n",
      "    b_c = 1.0/temp_array[ct]\n",
      "    b_n = 1.0/temp_array[newt]\n",
      "    if (-(b_n-b_c)*E+(f[newt]-f[ct])) >=0:\n",
      "        return newt\n",
      "    elif np.random.rand() < np.exp( (-(b_n-b_c)*E+(f[newt]-f[ct])) ):\n",
      "        return newt\n",
      "    else:\n",
      "        return ct\n",
      "    \n",
      "#Brownian Dynamics simulation\n",
      "def bd( x0, L, nex, kT_idx = None, kT = None, f = None):\n",
      "    dt =0.005\n",
      "    mass = 1\n",
      "    damping = 1\n",
      "    beta = 1.0/kT[kT_idx]\n",
      "    coeff_A = dt / ( mass * damping )\n",
      "    coeff_B = np.sqrt( 2.0 * dt / ( beta * mass * damping ) )\n",
      "    x = x0\n",
      "    P = [x0]\n",
      "    E = U( x0 )\n",
      "    T = [kT_idx]\n",
      "    u = [E]\n",
      "    for ex in xrange( nex ):\n",
      "        for t in xrange( L ):\n",
      "            g = G( x )\n",
      "            x = x - coeff_A * g + coeff_B * np.random.normal()\n",
      "            P.append( x )\n",
      "            T.append( kT_idx )\n",
      "            e = U( x )\n",
      "            u.append( e*beta )\n",
      "        kT_idx = change_temp( kT_idx, e, kT, f )\n",
      "        beta = 1.0/kT[kT_idx]\n",
      "        coeff_B = np.sqrt( 2.0 * dt / ( beta * mass * damping ) )\n",
      "    \n",
      "    return ( np.array( P, dtype=np.float64 ),np.array( T, dtype=np.intc ),np.array( u, dtype=np.float64 ).reshape(len(u),1) )\n",
      "\n",
      "# discretisation scheme\n",
      "def discretise( T, centers ):\n",
      "    D = np.zeros( shape=T.shape, dtype=np.int32 )\n",
      "    half_width = (centers[2]-centers[1])*0.5\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]+half_width:\n",
      "                    D[i] = j\n",
      "                    break\n",
      "    return D\n",
      "# generate trajectories\n",
      "def get_trajectories( centers, L, nex, kT_idx=0, kT = None ,f=None):\n",
      "    r = [] #trajectory list\n",
      "    #run the Brownian dynamics integrator\n",
      "    traj,therm, b = bd( -1.0, L ,nex, kT_idx = kT_idx, kT= kT, f=f)\n",
      "    dtraj = discretise( traj, centers ) #discretise the trajectory into the defined bins\n",
      "    r.append( { 'm':dtraj, 't': therm, 'b': b } ) # list of dictionaries in the correct pytram format\n",
      "    return r\n",
      "\n",
      "\n",
      "# set the discretisation and simulation parameters\n",
      "# kT array\n",
      "TEMPS = np.array([2.0,4.0,7.0, 15.0] )\n",
      "# factor used to make sure temperatures are roughly sampled equally\n",
      "f = np.array([-5.95994519, -3.37600304, -2.50314141, -2.04173675] ) \n",
      "#definition of the binning\n",
      "NBINS = 50\n",
      "centers = np.linspace(-0.6,4.2,NBINS)\n",
      "L = 100 #trajectoy length before an exchange is attempted   \n",
      "nex = 2000 #number of times a temperature exchange is attempted\n",
      "# run the simulations\n",
      "trajs = get_trajectories(centers, L, nex, kT_idx=0, kT=TEMPS, f=f)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "With the generated list of trajectories we can now use xTRAM to estimate the free energy for each bin for the lowest temperature we sampled at. This means, rather than just giving a list of trajectory dictionaries to the TRAMdata object, we also have to indicate the temperatures the trajectory was sampled at and the target temperature for which we want to know the final probability distribution/ free energy profile. "
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "tramdata = TRAMData(trajs, kT_K=TEMPS, kT_target=0)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Now we can just pass the TRAMdata object to the API function ```python xtram()``` indicate the lag at which the count matrix should be estimated and set the number of iterations and tolerance of the estimate. This will run the estimator."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "xtram_obj = xtram(tramdata, 1, maxiter=100, ftol=1.0e-15, verbose=True)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Here we compute the exact solution which can then be used to our estimate"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "Z = np.exp(-f)\n",
      "def prob( x, kT, Z ):\n",
      "    return exp(-U( x )/kT[0])/Z[0]\n",
      "exact = []\n",
      "for x in centers:\n",
      "    exact.append( prob( x, TEMPS, Z ) )\n",
      "exact = exact/np.sum(exact)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Now let us plot the results of the estimator and compare it to the exact solution. The central object from the estimator is xtram_obj.pi_i, which contains the stationary probability of each bin defined at the target temperature indicated. "
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "plt.figure(figsize=(8,6))\n",
      "plt.plot(centers,xtram_obj.f_i, color = 'b', linewidth =1.5, label ='xtram', marker = 'o')\n",
      "plt.plot(centers,-np.log(exact), color = 'k', linewidth = 1.5, label = 'exact')\n",
      "plt.plot(centers,-np.log(tramdata.N_K_i[0]/float(tramdata.N_K[0])), color = 'r', linestyle = '--', linewidth = 1.5, label = 'counts')\n",
      "plt.xlabel('x in [a.u.]', fontsize = 20)\n",
      "plt.ylabel('F in [kT]', fontsize = 20)\n",
      "plt.legend( loc =4)\n",
      "plt.tick_params(axis='both', which='major', labelsize=20)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "heading",
     "level": 2,
     "metadata": {},
     "source": [
      "Concluding remark"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "We hope that this example helps you to understand the basic usage of the xTRAM 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 xTRAM useful, please cite\\n\"\n",
      "xtram_obj.cite()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}