03-basin-hopping.ipynb

{
 "metadata": {
  "name": ""
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Basin-hopping with leonard jones clusters",
      "Implementing Wales and Dolye's [paper](http://pubs.acs.org/doi/abs/10.1021/jp970984n) in OpenMM."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from math import ceil\n",
      "import matplotlib.pyplot as pp\n",
      "from itertools import islice\n",
      "import numpy as np\n",
      "from simtk.unit import angstroms\n",
      "import simtk.openmm as mm"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 1
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "N_PARTICLES=23\n",
      "system = mm.System()\n",
      "force = mm.CustomNonbondedForce('4*(r6*r6 - r6); r6 = 1/(r*r*r*r*r*r);')\n",
      "force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)\n",
      "system.setDefaultPeriodicBoxVectors((10,0,0), (0,10,0), (0,0,10))\n",
      "force.setCutoffDistance(5)\n",
      "\n",
      "for i in range(N_PARTICLES):\n",
      "    system.addParticle(1)\n",
      "    force.addParticle([])\n",
      "system.addForce(force)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 2,
       "text": [
        "0"
       ]
      }
     ],
     "prompt_number": 2
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "integrator = mm.CustomIntegrator(0.001)\n",
      "integrator.addGlobalVariable('a_cur', 0)\n",
      "integrator.addGlobalVariable('a_old', 0)\n",
      "integrator.addPerDofVariable('y_cur', 0)\n",
      "integrator.addPerDofVariable('y_old', 0)\n",
      "integrator.addComputeGlobal('a_cur', '0.5*(1+sqrt(1+(4*a_old*a_old)))')\n",
      "integrator.addComputeGlobal('a_old', 'a_cur')\n",
      "integrator.addComputePerDof('y_cur', 'x + dt*f')\n",
      "integrator.addComputePerDof('y_old', 'y_cur')\n",
      "integrator.addComputePerDof('x', 'y_cur + (a_old - 1) / a_cur * (y_cur - y_old)')\n",
      "\n",
      "context = mm.Context(system, integrator)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 3
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def grid_coordinates(n_points):\n",
      "    \"\"\"Coordinates of grid points in 3d (e.g. (0,0,0), (0,0,1), ...)\n",
      "\n",
      "    Returns\n",
      "    -------\n",
      "    points : shape=(n_points, 3)\n",
      "    \"\"\"\n",
      "    griddim = ceil(n_points**(1.0/3.0))\n",
      "    mesh = np.indices((griddim, griddim, griddim))\n",
      "\n",
      "    points = np.array(list(zip(*[e.flatten() for e in mesh])))[:N_PARTICLES]\n",
      "    return points\n",
      "\n",
      "def scatter3d(context):\n",
      "    xyz = context.getState(getPositions=True, enforcePeriodicBox=True).getPositions(asNumpy=True)._value\n",
      "    fig = pp.figure()\n",
      "    ax = fig.add_subplot(111, projection='3d')\n",
      "    ax.scatter(xyz[:,0], xyz[:,1], xyz[:,2])\n",
      "    return ax"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 4
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def basin_hop(context, scaling=0.3, temperature=0.8, n_steps_per_hop=100, n_hops=1000):\n",
      "    state = context.getState(getPositions=True, getEnergy=True, enforcePeriodicBox=True)\n",
      "    positions = state.getPositions(asNumpy=True)._value\n",
      "    energy = state.getPotentialEnergy()._value\n",
      "    n_acceptances = 0\n",
      "\n",
      "    best_energy = 0\n",
      "    best_positions = None\n",
      "    \n",
      "    for i in range(n_hops):\n",
      "        displacement = scaling * np.random.random(positions.shape)\n",
      "        context.setPositions(positions + displacement)\n",
      "        context.getIntegrator().step(n_steps_per_hop)\n",
      "        \n",
      "        new_state = context.getState(getEnergy=True, getPositions=True, enforcePeriodicBox=True)\n",
      "        new_energy = new_state.getPotentialEnergy()._value\n",
      "        p = np.exp(-(new_energy - energy) / temperature)\n",
      "        \n",
      "        if (np.random.rand() < p):\n",
      "            energy = new_energy\n",
      "            positions = new_state.getPositions(asNumpy=True)._value\n",
      "            n_acceptances += 1\n",
      "        if i % 1000 == 0 and i > 0:\n",
      "            print('step=%d, energy=%.3f, MC acceptance=%.3f' % (i, energy, float(n_acceptances) / i))\n",
      "\n",
      "        if energy < best_energy:\n",
      "            best_positions = positions\n",
      "\n",
      "    context.setPositions(best_positions)\n",
      "    mm.LocalEnergyMinimizer.minimize(context, 1e-6)\n",
      "\n",
      "    print('final energy', context.getState(getEnergy=True).getPotentialEnergy()._value)\n",
      "    print('Acceptance: %d/%d = %.3f' % (n_acceptances, n_hops, float(n_acceptances) / n_hops))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 5
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "context.setPositions(grid_coordinates(N_PARTICLES))\n",
      "basin_hop(context, n_hops=10000)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "step=1000, energy=-87.443, MC acceptance=0.511\n",
        "step=2000, energy=-87.432, MC acceptance=0.504"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        "step=3000, energy=-87.405, MC acceptance=0.507"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        "step=4000, energy=-87.436, MC acceptance=0.507"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        "step=5000, energy=-87.418, MC acceptance=0.506"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        "step=6000, energy=-87.428, MC acceptance=0.502"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        "step=7000, energy=-87.234, MC acceptance=0.506"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        "step=8000, energy=-87.415, MC acceptance=0.511"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        "step=9000, energy=-87.443, MC acceptance=0.516"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        "('final energy', -87.44991302490234)"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        "Acceptance: 5184/10000 = 0.518\n"
       ]
      }
     ],
     "prompt_number": 6
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from simtk.openmm.app.element import carbon\n",
      "from simtk.openmm.app import Topology, PDBFile\n",
      "topology = Topology()\n",
      "chain = topology.addChain()\n",
      "residue = topology.addResidue('A', chain)\n",
      "for i in range(N_PARTICLES):\n",
      "    topology.addAtom('XX', carbon, residue)\n",
      "a,b,c = system.getDefaultPeriodicBoxVectors()\n",
      "topology.setUnitCellDimensions((a[0]._value, b[1]._value, c[2]._value) * angstroms)\n",
      "\n",
      "with open('out.pdb', 'w') as f:\n",
      "    PDBFile.writeFile(topology, context.getState(getPositions=True, enforcePeriodicBox=True).getPositions()._value, f)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 7
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "print context.getState(getPositions=True, enforcePeriodicBox=True).getPositions(asNumpy=True)._value"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "[[ 9.91630173  0.36938968  7.5276351 ]\n",
        " [ 9.20100689  9.76116753  8.10398006]\n",
        " [ 8.69661045  9.06374073  8.83127594]\n",
        " [ 0.06466249  0.26453698  8.64702129]\n",
        " [ 9.12388325  0.03974746  9.18732071]\n",
        " [ 9.26640606  9.0619421   9.78249073]\n",
        " [ 9.80602646  0.8213762   9.58628273]\n",
        " [ 9.2463789   0.07972761  0.27695048]\n",
        " [ 8.32210445  9.59865093  9.80305862]\n",
        " [ 0.30238992  9.43979073  7.96023226]\n",
        " [ 9.42557144  8.66699409  8.02242184]\n",
        " [ 8.795228    9.19381237  0.76314068]\n",
        " [ 9.59261322  9.33237839  7.14331436]\n",
        " [ 9.77092361  9.24324322  8.86161137]\n",
        " [ 0.34914479  8.49871063  7.41846085]\n",
        " [ 0.88608289  0.493341    9.33897781]\n",
        " [ 0.12986501  9.79016304  9.70449162]\n",
        " [ 9.93072796  9.22561741  0.62665796]\n",
        " [ 1.25692058  8.90738487  8.00583839]\n",
        " [ 0.41075778  8.44400597  8.53698444]\n",
        " [ 9.44068909  8.22970963  9.06386566]\n",
        " [ 0.83877373  9.43383217  8.9269228 ]\n",
        " [ 0.31151932  8.70152855  9.64017868]]\n"
       ]
      }
     ],
     "prompt_number": 9
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}