VPythonAnimations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Taylan Cemgil\n",
    "### Boğaziçi University, Dept. of Computer Engineering\n",
    "\n",
    "The animations on this notebook need vpython (http://vpython.org/) installed on your system.  For some peculiar reason, the kernel needs to be restarted after every run."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ball jumping around a circle\n",
    "<img src=\"images/circular-robot.png\" width=\"500\" align=\"right\">\n",
    "\n",
    "A toy example for localization using an Hidden Markov Model. The ball stays put with probability $\\pi_0$, jumps with probability $\\pi_1$ forward, and is 'kidnapped' with a small probability $\\pi_2$. At each step, we observe the color of the current tile.\n",
    "\n",
    "At every step, the filtering density over the possible locations are shown.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "require.undef(\"nbextensions/vpython_libraries/glow.min\");"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "require.undef(\"nbextensions/vpython_libraries/glowcomm\");"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "require.undef(\"nbextensions/vpython_libraries/jquery-ui.custom.min\");"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "require([\"nbextensions/vpython_libraries/glow.min\"], function(){console.log(\"GLOW LOADED\");})"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "require([\"nbextensions/vpython_libraries/glowcomm\"], function(){console.log(\"GLOWCOMM LOADED\");})"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "require([\"nbextensions/vpython_libraries/jquery-ui.custom.min\"], function(){console.log(\"JQUERY LOADED\");})"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div id=\"glowscript\" class=\"glowscript\"></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "window.__context = { glowscript_container: $(\"#glowscript\").removeAttr(\"id\")}"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-d6ca11575286>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m    102\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    103\u001b[0m \u001b[0;32mwhile\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 104\u001b[0;31m     \u001b[0mvp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrate\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mRATE\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    105\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    106\u001b[0m     \u001b[0mYellowBall\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpos\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mYellowBall\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpos\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mYellowBall\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvel\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mdt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Users/cemgil/anaconda/lib/python3.6/site-packages/vpython/vpython.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, N)\u001b[0m\n\u001b[1;32m    201\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrval\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mN\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    202\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrval\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"rate value must be greater than or equal to 1\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 203\u001b[0;31m         \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_RateKeeper2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrval\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m## calls __call__ in rate_control.py\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    204\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    205\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0msys\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mversion\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;34m'3'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Users/cemgil/anaconda/lib/python3.6/site-packages/vpython/rate_control.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, maxRate)\u001b[0m\n\u001b[1;32m    205\u001b[0m             \u001b[0mdt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlastSleep\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcalls\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muserTime\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcallTime\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdelay\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m\\\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    206\u001b[0m                  \u001b[0mrenders\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrenderTime\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0msleeps\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minteractionPeriod\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0m_clock\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 207\u001b[0;31m             \u001b[0m_sleep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    208\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlastSleep\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_clock\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    209\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcalls\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Users/cemgil/anaconda/lib/python3.6/site-packages/vpython/rate_control.py\u001b[0m in \u001b[0;36m_sleep\u001b[0;34m(dt)\u001b[0m\n\u001b[1;32m     47\u001b[0m         \u001b[0mdtsleep\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnticks\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0m_tick\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     48\u001b[0m         \u001b[0mt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_clock\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 49\u001b[0;31m         \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msleep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdtsleep\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     50\u001b[0m         \u001b[0mt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_clock\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     51\u001b[0m         \u001b[0mdt\u001b[0m \u001b[0;34m-=\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "from __future__ import division, print_function\n",
    "#from math import *\n",
    "import vpython as vp\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pylab as plt\n",
    "import matplotlib as mpl\n",
    "\n",
    "from itertools import product\n",
    "\n",
    "# Simulation parameters\n",
    "RATE = 150\n",
    "dt = 1./RATE\n",
    "g_earth = 19.8\n",
    "T_period = 0.5\n",
    "\n",
    "showFilteringDensity = True\n",
    "\n",
    "# Set the scene\n",
    "\n",
    "# Floor\n",
    "FloorLenght = 7\n",
    "FloorHeight = 7\n",
    "floorYPos = -1.5\n",
    "floor = vp.box(pos=vp.vector(0,floorYPos,0), length=FloorLenght, height=0.05, width=FloorHeight, color=vp.vector(0.4,0.4,0.4))\n",
    "\n",
    "# Number of Tiles\n",
    "L = 40\n",
    "tile_radius = 0.1\n",
    "TileCoordinate = [(2*(np.cos(th)+np.sin(th)), 0, 2*(np.cos(th)-np.sin(th)) ) for th in np.linspace(0,2*np.pi,L)]\n",
    "\n",
    "ColorOfObservation = [vp.color.white, vp.color.blue]\n",
    "TileColor = []\n",
    "FilteringDensityCylinder = []\n",
    "for i in range(L):\n",
    "    sx, sy, sz = TileCoordinate[i]\n",
    "    s = vp.vector(sx,sy,sz)\n",
    "    # Pick a random color for the tile \n",
    "    c = np.random.choice(range(len(ColorOfObservation)))\n",
    "    TileColor.append(c)\n",
    "    vp.cylinder(pos=s, axis=vp.vector(0,-tile_radius,0), color=ColorOfObservation[c], radius=tile_radius)\n",
    "    \n",
    "    if showFilteringDensity:\n",
    "        s2 = vp.vector(sx,sy+floorYPos-0.1,sz)\n",
    "        cyl = vp.cylinder(pos=s2, axis=vp.vector(0,0.4,0), radius=tile_radius)\n",
    "        FilteringDensityCylinder.append(cyl)\n",
    "\n",
    "        \n",
    "def Setup_HMM_parameters():\n",
    "    # Probability of staying on the same tile\n",
    "    ep = 0.2\n",
    "    # Probability of making an arbitrary jump\n",
    "    kidnap = 0.01\n",
    "    # Probability of correctly observing the tile color\n",
    "    a = 0.99\n",
    "\n",
    "    # Set up the transition matrix\n",
    "    idx = [i for i in range(1,L)]+[0]\n",
    "    I = np.diag(np.ones(L))\n",
    "    A = (1-kidnap)*(ep*I + (1-ep)*I[:,idx]) + kidnap*np.ones((L,L))/L\n",
    "    C = np.zeros((2,L))\n",
    "    pred = np.ones(L)/L\n",
    "\n",
    "    for i in range(L):\n",
    "        C[0,i] = a*(1 - TileColor[i]) + (1-a)*TileColor[i]\n",
    "        C[1,i] = a*TileColor[i] + (1-a)*(1 - TileColor[i])\n",
    "\n",
    "    return A, C, pred\n",
    "\n",
    "A, C, pred = Setup_HMM_parameters()\n",
    "    \n",
    "    \n",
    "Obs_noiseless = vp.sphere(pos=vp.vector(0,floorYPos-0.2,0), color=vp.color.black, radius=tile_radius)\n",
    "Obs = vp.sphere(pos=vp.vector(0,floorYPos+0.2,0), color=vp.color.black, radius=tile_radius)\n",
    "\n",
    "\n",
    "\n",
    "nf = mpl.colors.Normalize(vmin=0, vmax=0.7, clip=True)\n",
    "cmap = plt.cm.ScalarMappable(cmap=plt.cm.hot, norm=nf)\n",
    "lamp = vp.local_light(pos=vp.vector(5,-4,0), color=vp.color.white)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# Select a random initial state\n",
    "CurrentState = np.random.choice(range(L))\n",
    "sx, sy, sz = TileCoordinate[CurrentState]\n",
    "\n",
    "# Set the initial state of the yellow ball\n",
    "x0 = vp.vector(sx,sy,sz)\n",
    "v0 = vp.vector(0,0,0)\n",
    "YellowBall = vp.sphere(pos=x0, radius=tile_radius, color=vp.color.yellow,\n",
    "             make_trail=True, interval=1, retain=RATE*T_period)\n",
    "YellowBall.vel = v0\n",
    "YellowBall.g = g_earth \n",
    "YellowBall.T_period = T_period \n",
    "YellowBall.retain = RATE*YellowBall.T_period\n",
    "\n",
    "FlightCounter = YellowBall.T_period/dt\n",
    "    \n",
    "\n",
    "\n",
    "while 1:\n",
    "    vp.rate (RATE)\n",
    "\n",
    "    YellowBall.pos = YellowBall.pos + YellowBall.vel*dt\n",
    "    \n",
    "    # If hit halfway, dim the observation\n",
    "    if FlightCounter>= YellowBall.T_period/dt/2.:\n",
    "        Obs.color = vp.color.black\n",
    "        Obs_noiseless.color = vp.color.black\n",
    "\n",
    "    # If arrived, dim the observation\n",
    "    if FlightCounter>= YellowBall.T_period/dt:\n",
    "        observation = np.random.choice(range(C.shape[0]), p=C[:,CurrentState])\n",
    "        o_noiseless = TileColor[CurrentState]\n",
    "        Obs.color = ColorOfObservation[observation]\n",
    "        Obs_noiseless.color = ColorOfObservation[o_noiseless]\n",
    "        \n",
    "        \n",
    "        NextState = np.random.choice(range(A.shape[0]), p=A[:,CurrentState])\n",
    "\n",
    "        pred = C[observation,:]*pred\n",
    "        pred = pred/np.sum(pred)\n",
    "\n",
    "        if showFilteringDensity:\n",
    "            for k in range(L):\n",
    "                col = cmap.to_rgba(pred[k])\n",
    "                vcol = vp.vector(col[0], col[1], col[2])\n",
    "                FilteringDensityCylinder[k].color = vcol\n",
    "                FilteringDensityCylinder[k].axis=vp.vector(0,pred[k]+0.15,0)\n",
    "\n",
    "        pred = A.dot(pred)\n",
    "\n",
    "\n",
    "        ## Plan the jump\n",
    "        sx, sy, sz = TileCoordinate[CurrentState]\n",
    "        tx, ty, tz = TileCoordinate[NextState]\n",
    "        v_vert = YellowBall.g*YellowBall.T_period/2 + (ty-sy)/YellowBall.T_period\n",
    "\n",
    "        YellowBall.vel = vp.vector((tx-sx)/YellowBall.T_period, v_vert,(tz-sz)/YellowBall.T_period)   \n",
    "        YellowBall.pos = vp.vector(sx,sy,sz)\n",
    "        #\n",
    "        CurrentState = NextState\n",
    "        FlightCounter = 0\n",
    "    else:\n",
    "        YellowBall.vel.y = YellowBall.vel.y - YellowBall.g*dt\n",
    "        FlightCounter +=1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Proposal Mechanism of Metropolis Hastings\n",
    "<img src=\"images/mh-chain.png\" width=\"500\" align=\"right\">\n",
    "\n",
    "This demo illustrates a MH chain on a discrete state space. The size of each tile $x$ is proportional to the target probability $\\pi(x)$.\n",
    "\n",
    "The yellow big ball shows the current state $x$ of the chain, and the small gray ball shows the proposed state $x'$ drawn from a proposal density $q(\\cdot|x)$. The proposed new state is accepted with probability \n",
    "$$\n",
    "\\min\\left\\{1, \\frac{\\pi(x') q(x|x')}{\\pi(x) q(x'|x)} \\right\\}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "require.undef(\"nbextensions/vpython_libraries/glow.min\");"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "require.undef(\"nbextensions/vpython_libraries/glowcomm\");"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "require.undef(\"nbextensions/vpython_libraries/jquery-ui.custom.min\");"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "require([\"nbextensions/vpython_libraries/glow.min\"], function(){console.log(\"GLOW LOADED\");})"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "require([\"nbextensions/vpython_libraries/glowcomm\"], function(){console.log(\"GLOWCOMM LOADED\");})"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "require([\"nbextensions/vpython_libraries/jquery-ui.custom.min\"], function(){console.log(\"JQUERY LOADED\");})"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div id=\"glowscript\" class=\"glowscript\"></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "window.__context = { glowscript_container: $(\"#glowscript\").removeAttr(\"id\")}"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-cd9b9eeab396>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m    120\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    121\u001b[0m \u001b[0;32mwhile\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 122\u001b[0;31m     \u001b[0mvp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrate\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mRATE\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    123\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    124\u001b[0m     \u001b[0mB\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mpQ\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpos\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mB\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mpQ\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpos\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mB\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mpQ\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvel\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mdt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Users/cemgil/anaconda/envs/py35/lib/python3.5/site-packages/vpython/vpython.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, N)\u001b[0m\n\u001b[1;32m    201\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrval\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mN\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    202\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrval\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"rate value must be greater than or equal to 1\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 203\u001b[0;31m         \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_RateKeeper2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrval\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m## calls __call__ in rate_control.py\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    204\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    205\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0msys\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mversion\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;34m'3'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Users/cemgil/anaconda/envs/py35/lib/python3.5/site-packages/vpython/rate_control.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, maxRate)\u001b[0m\n\u001b[1;32m    205\u001b[0m             \u001b[0mdt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlastSleep\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcalls\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muserTime\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcallTime\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdelay\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m\\\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    206\u001b[0m                  \u001b[0mrenders\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrenderTime\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0msleeps\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minteractionPeriod\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0m_clock\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 207\u001b[0;31m             \u001b[0m_sleep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    208\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlastSleep\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_clock\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    209\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcalls\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Users/cemgil/anaconda/envs/py35/lib/python3.5/site-packages/vpython/rate_control.py\u001b[0m in \u001b[0;36m_sleep\u001b[0;34m(dt)\u001b[0m\n\u001b[1;32m     47\u001b[0m         \u001b[0mdtsleep\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnticks\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0m_tick\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     48\u001b[0m         \u001b[0mt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_clock\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 49\u001b[0;31m         \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msleep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdtsleep\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     50\u001b[0m         \u001b[0mt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_clock\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     51\u001b[0m         \u001b[0mdt\u001b[0m \u001b[0;34m-=\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "from __future__ import division, print_function\n",
    "from math import *\n",
    "import vpython as vp\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pylab as plt\n",
    "import matplotlib as mpl\n",
    "\n",
    "from vpython_utilities import make_grid2D\n",
    "\n",
    "# Simulation speed\n",
    "RATE = 300\n",
    "dt = 1./RATE\n",
    "\n",
    "W1 = 1\n",
    "W2 = 6\n",
    "step = 0.5\n",
    "\n",
    "n2 = int(W2/step)+1\n",
    "n1 = int(W1/step)+1\n",
    "PointList, sub2ind, ind2sub, edges, A = make_grid2D(n2,n1)\n",
    "\n",
    "Trans = A/np.sum(A, axis=0, keepdims=True)\n",
    "Trans = Trans.dot(Trans).dot(Trans)\n",
    "L = len(PointList)\n",
    "\n",
    "Y = []\n",
    "for i in range(L):\n",
    "    p = PointList[i]\n",
    "    x = p[0]*step-W2/2.\n",
    "    z = p[2]*step-W1/2.\n",
    "    E = 2+ np.cos(2*pi*z/3)+np.sin(2*pi*x/5)+np.random.randn()/10.\n",
    "    y = 2*np.exp(-1.1*E)\n",
    "    PointList[i] = (x, 0, z)\n",
    "    Y.append(y)\n",
    "    \n",
    "MAX = 1\n",
    "MIN = 0\n",
    "nf = mpl.colors.Normalize(vmin=MIN, vmax=MAX, clip=True)\n",
    "cmap = plt.cm.ScalarMappable(cmap=plt.cm.cool_r, norm=nf)\n",
    "\n",
    "#floor = box(pos=vector(0,-0.04,0), length=W1, height=0.05, width=W2, color=color.black)\n",
    "\n",
    "wd = 0.4\n",
    "radius = 0.2\n",
    "maxY = max(Y)\n",
    "\n",
    "for i in range(L):\n",
    "    sx, sy, sz = PointList[i]\n",
    "    s = vp.vector(sx,-radius,sz)\n",
    "    #vp.sphere(pos=s, color=vp.color.cyan, radius=0.1)\n",
    "    #vcol = vp.vector(0.9,0.9,0.9)\n",
    "    col = cmap.to_rgba(Y[i]/maxY)\n",
    "    vcol = vp.vector(col[0], col[1], col[2])\n",
    "    #vp.cylinder(pos=s, axis=vp.vector(0,-0.1,0), color=vcol, radius=radius*np.sqrt(nf(sy))) \n",
    "    \n",
    "    wd = step*np.sqrt(Y[i]/maxY)\n",
    "    vp.box(pos=s,length=wd, height=0.05, width=wd, color=vcol )\n",
    "    #s = vp.vector(sx,(sy-radius)/2.,sz)\n",
    "    #vp.box(pos=s,length=wd, height=sy-radius, width=wd, color=vcol )\n",
    "\n",
    "Cur = []\n",
    "\n",
    "# Cnt[i] is the number of ticks after a new movement has started of the i'th particle\n",
    "Cnt = []\n",
    "\n",
    "B = []\n",
    "\n",
    "\n",
    "g_earth = 49.8\n",
    "T_period = 0.25\n",
    "\n",
    "N = 2\n",
    "\n",
    "for i in range(N):  \n",
    "    cur = np.random.choice(range(L))\n",
    "    sx, sy, sz = PointList[cur]\n",
    "    x0 = vp.vector(sx,sy,sz)\n",
    "    v0 = vp.vector(0,0,0)\n",
    "    ball = vp.sphere(pos=x0, radius=radius, color=vp.color.yellow,\n",
    "                 make_trail=True, interval=1, retain=RATE*T_period)\n",
    "    ball.vel = v0\n",
    "    ball.g = g_earth\n",
    "    ball.T_period = T_period \n",
    "    ball.retain = RATE*ball.T_period\n",
    "    cnt = ball.T_period/dt\n",
    "    B.append(ball)\n",
    "    Cur.append(cur)\n",
    "    Cnt.append(cnt)\n",
    "\n",
    "lamp = vp.local_light(pos=vp.vector(0,-1,0), color=vp.color.yellow)\n",
    "\n",
    "    \n",
    "def selectNextState(cur):\n",
    "    pr = Trans[:,cur]\n",
    "    nex = np.random.choice(range(L), p=pr) \n",
    "    lw = np.log(Trans[cur, nex]) - np.log(Trans[nex, cur])\n",
    "    return nex, lw\n",
    "\n",
    "def planJump(ball, curPos, nexPos):\n",
    "    sx, sy, sz = curPos\n",
    "    tx, ty, tz = nexPos\n",
    "    v_vert = ball.g*ball.T_period/2 + (ty-sy)/ball.T_period                \n",
    "    vel = vp.vector((tx-sx)/ball.T_period, v_vert,(tz-sz)/ball.T_period)   \n",
    "    pos = vp.vector(sx,sy,sz)  \n",
    "    return pos, vel\n",
    "\n",
    "\n",
    "# Particle index of the Chain\n",
    "pP = 0\n",
    "# Particle index of the proposal\n",
    "pQ = 1\n",
    "B[pQ].make_trail = False\n",
    "B[pQ].color = vp.vector(0.6,0.6,0.6)\n",
    "B[pQ].radius = radius/2\n",
    "\n",
    "# Is proposal ball moving?\n",
    "pQmove = True\n",
    "log_q_ratio = 0\n",
    "\n",
    "while 1:\n",
    "    vp.rate (RATE)\n",
    "    \n",
    "    B[pQ].pos = B[pQ].pos + B[pQ].vel*dt\n",
    "    B[pP].pos = B[pP].pos + B[pP].vel*dt\n",
    "    if Cnt[pQ]>= B[pQ].T_period/dt:\n",
    "        \n",
    "        if pQmove:\n",
    "            accept = np.log(np.random.rand()) < log_q_ratio + np.log(Y[Cur[pQ]]) - np.log(Y[Cur[pP]])\n",
    "\n",
    "            if accept:\n",
    "                \n",
    "                # pP jumps to new location\n",
    "                B[pP].g = g_earth\n",
    "                pos, vel = planJump(B[pP], PointList[Cur[pP]], PointList[Cur[pQ]])              \n",
    "                B[pP].vel = vel \n",
    "                B[pP].pos = pos\n",
    "                \n",
    "                # pQ stays put\n",
    "                B[pQ].vel = vp.vector(0,0,0)\n",
    "                B[pQ].pos.x, B[pQ].pos.y, B[pQ].pos.z = PointList[Cur[pQ]]\n",
    "                B[pQ].g  = 0\n",
    "\n",
    "                Cur[pP] = Cur[pQ]\n",
    "            else:\n",
    "                # pP jumps vertically\n",
    "                B[pP].g = g_earth\n",
    "                pos, vel = planJump(B[pP], PointList[Cur[pP]], PointList[Cur[pP]])  \n",
    "                B[pP].vel = vel\n",
    "                B[pP].pos = pos\n",
    "\n",
    "\n",
    "                # pQ disappears\n",
    "                B[pQ].visible = False\n",
    "                B[pQ].g = g_earth/10.\n",
    "                pos, vel = planJump(B[pQ], PointList[Cur[pQ]], PointList[Cur[pP]])              \n",
    "                B[pQ].vel = vel \n",
    "                B[pQ].pos = pos\n",
    "\n",
    "                Cur[pQ] = Cur[pP]               \n",
    "            pQmove = False\n",
    "        else:\n",
    "            B[pQ].visible = True\n",
    "\n",
    "            nex, log_q_ratio = selectNextState(Cur[pP])\n",
    "            # pP stays put\n",
    "            B[pP].vel = vp.vector(0,0,0)\n",
    "            B[pP].pos.x, B[pP].pos.y, B[pP].pos.z = PointList[Cur[pP]]\n",
    "            B[pP].g  = 0\n",
    "\n",
    "            # pQ jumps to new location\n",
    "            B[pQ].g = g_earth/10.\n",
    "            pos, vel = planJump(B[pQ], PointList[Cur[pP]], PointList[nex])              \n",
    "            B[pQ].vel = vel \n",
    "            B[pQ].pos = pos\n",
    "\n",
    "            Cur[pQ] = nex               \n",
    "            \n",
    "            pQmove = True\n",
    "\n",
    "                \n",
    "        Cnt[pP] = 0\n",
    "        Cnt[pQ] = 0\n",
    "                \n",
    "    else:\n",
    "        B[pP].vel.y = B[pP].vel.y - B[pP].g*dt\n",
    "        B[pQ].vel.y = B[pQ].vel.y - B[pQ].g*dt\n",
    "        Cnt[pP] +=1\n",
    "        Cnt[pQ] +=1\n",
    "         "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [default]",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": "block",
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}