Lecture12.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lecture 12: Simple Random Systems\n",
    "\n",
    "**Overview:**\n",
    "* Random walk in 1D.\n",
    "* Random walk in 2D and entropy.\n",
    "\n",
    "**Next Lecture:**\n",
    "* Intro to Monte Carlo, integrating functions in high dimensions.\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tasks\n",
    "* Modify randwalk() to use purely numpy arrays i.e. *remove the loop over nsamples*. Hint: np.random.randint() can generate array's of arbitrary shape. What sort of performance gain do you achieve with this?\n",
    "\n",
    "* Modify randwalk() to take random step sizes (between 0 and 1) does the slope of $<x^2>$ vs $n$ change?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook \n",
    "# change all of these to inline\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# we will use numba again\n",
    "from numba import jit\n",
    "import matplotlib.cm as cm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def randwalk(nsteps = 100, nsamples = 1000):\n",
    "    nsteps = int(nsteps)\n",
    "    nsamples = int(nsamples)\n",
    "    rms = np.zeros(nsteps)\n",
    "    trajectories = np.zeros([nsamples, nsteps])\n",
    "    \n",
    "    for ii in range(nsamples):\n",
    "        # draw a random integer between -1 and 1 for each step, and scale to -0.5 - 0.5\n",
    "        # exclude first point, becuase we want to start at 0\n",
    "        \n",
    "        x0 = 2*np.random.randint(0,1+1, nsteps-1) -1  \n",
    "        \n",
    "        #start at 0\n",
    "        x0 = np.insert(x0,0,0,axis= 0)\n",
    "        \n",
    "        # take steps\n",
    "        traj = np.cumsum(x0)\n",
    "        \n",
    "        # update trajectories array and displacement squared array\n",
    "        trajectories[ii,:] = traj\n",
    "        rms+=traj**2\n",
    "    \n",
    "    trajectories = np.array(trajectories)\n",
    "    return rms/nsamples, trajectories\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def randwalk_np(nsteps = 100, nsamples = 1000):\n",
    "    nsteps = int(nsteps)\n",
    "    nsamples = int(nsamples)\n",
    "       \n",
    "    # you need at most 3 more lines of code here\n",
    "\n",
    "    trajectories = []\n",
    "    rms = np.mean(trajectories**2, axis = 0)\n",
    "    return rms, trajectories\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%%timeit   # run some timing routines on the current cell, only uncomment for testing purposes, se\n",
    "nsteps = 500\n",
    "nsamples = 2e4\n",
    "rms, trajectories = randwalk(nsteps, nsamples)\n",
    "#rms, trajectories = randwalk_np(nsteps, nsamples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "# plot every 500th trajectory, so the plot is not too crowded\n",
    "for ii in np.arange(0,len(rms)-1,400):\n",
    "    ax.plot(trajectories[ii,:],'.-')\n",
    "\n",
    "# put a horizontal line at y=0\n",
    "ax.hlines(0,0,1.1*nsteps, linestyles = 'dashed', zorder = 10) #zorder sets the layer order in the plot\n",
    "\n",
    "ax.set_xlim([0,nsteps])\n",
    "ax.set_xlabel('Step number', fontsize =14)\n",
    "ax.set_ylabel('$x$', fontsize = 14)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examine the statistics: mean = $<x>$and variance =  $<x^2>$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax2 = fig.add_subplot(121)\n",
    "ax1 = fig.add_subplot(122)\n",
    "steps = np.arange(len(rms))\n",
    "\n",
    "# function to make a straight line\n",
    "# lambda is a special constructor to create functions on the fly\n",
    "linear = lambda D, x : 2*D*x\n",
    "\n",
    "ax1.plot(rms, 'o')\n",
    "ax1.plot(steps, linear(0.5,steps), lw = 2)\n",
    "ax1.set_xlabel('Step number', fontsize =14)\n",
    "ax1.set_ylabel('$<x^2>$', fontsize = 14)\n",
    "\n",
    "ax2.plot(np.mean(trajectories,axis = 0), '-')\n",
    "ax2.set_ylim([-1,1])\n",
    "ax2.set_xlabel('Step number', fontsize =14)\n",
    "ax2.set_ylabel('$<x>$', fontsize = 14)\n",
    "fig.subplots_adjust(left = 0.12,right = 0.96, bottom = 0.14, top = 0.86,wspace =0.45, hspace = 0.35)\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examine the probability distributions for 1D random walk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's look at the distribution of x for 3 different timepoints\n",
    "\n",
    "# define a Gaussin function for plotting, use the lambda constructor again\n",
    "Gaussian = lambda Dt, x : 1/np.sqrt(4*np.pi*Dt)*np.exp(-x**2/(4*Dt))\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "# plot P(x) for three differnt timepoints\n",
    "for ts in [50,100,200]:\n",
    "    Dt = 0.5*ts\n",
    "    \n",
    "    # P(x) is obtained from the histogram over all trajectories at a fixed time point\n",
    "    n, bins = np.histogram(trajectories[:,ts], bins = 'auto', density = True) \n",
    "    \n",
    "    # histogram returns bin edges, need to extract bin centers for plotting purposes\n",
    "    bin_center = np.diff(bins)+bins[:-1]\n",
    "\n",
    "    sf = np.max(n)*np.sqrt(4*np.pi*Dt) # correct scaling for histogram normalization\n",
    "\n",
    "    # make a bar chart\n",
    "    ax.bar(bin_center,n)\n",
    "    \n",
    "    # also plot points\n",
    "    # assign the plot object to a variable so we can 'grab' its color\n",
    "    points, = ax.plot(bin_center,n ,'o', mfc = 'white', label = 't = {} steps'.format(ts))\n",
    "    \n",
    "    # plot exact Gaussian with sigma = sqrt(2Dt)\n",
    "    ax.plot(bin_center, sf*Gaussian(Dt, bin_center), lw = 2, color = points.get_color())\n",
    "\n",
    "ax.set_ylabel('Probablity')\n",
    "ax.set_xlabel('x')\n",
    "ax.legend()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Random walk in Two-Dimensions\n",
    "* using numba jit to speed up critical components, there are many loops and simple conditionals here, numba can result in a significant speed up.\n",
    "* make sure you understand to flow of the code. What happens if you remove the `@jit` decorator from `step_2d`?\n",
    "\n",
    "* try commenting out the jit decorator around step_2d and see how fast the code runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def step_2d():\n",
    "    # step for random walk on 2D square lattice\n",
    "    d = np.random.randint(1,5)\n",
    "    if d == 1:\n",
    "        x_s, y_s = 1, 0\n",
    "    if d == 2:\n",
    "        x_s, y_s = 0, 1\n",
    "    if d == 3:\n",
    "        x_s, y_s = -1, 0\n",
    "    if d == 4:\n",
    "        x_s, y_s = 0, -1\n",
    "    return [x_s, y_s]\n",
    "\n",
    "\n",
    "@jit(nopython=True)\n",
    "def randwalk_2d(nsteps, wall = 1e6):\n",
    "   # trajectory for a single 2D random walk\n",
    "    x = np.zeros(nsteps+1, np.double)\n",
    "    y = np.zeros(nsteps+1, np.double)\n",
    "    for i in range(nsteps):\n",
    "        s = step_2d()\n",
    "        x[i+1] = x[i] + s[0]\n",
    "        if abs(x[i+1]) > wall: x[i+1] = x[i]\n",
    "        y[i+1] = y[i] + s[1]\n",
    "        if abs(y[i+1]) > wall: y[i+1] = y[i]\n",
    "\n",
    "    return x, y \n",
    "\n",
    "@jit(nopython=True)\n",
    "def ensemble_2d(N,t_max, wall = 1e6):\n",
    "    #generate an ensemble or random walkers\n",
    "    xh = np.zeros(N, np.double)\n",
    "    yh = np.zeros(N, np.double)\n",
    "\n",
    "    for i in range(N):\n",
    "        x, y = randwalk_2d(t_max, wall)\n",
    "        xh[i] = x[-1]\n",
    "        yh[i] = y[-1]\n",
    "\n",
    "    return xh, yh\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 5000\n",
    "x,y = randwalk_2d(n, wall = 20)\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "colormap = cm.Reds\n",
    "#ax.scatter(x, y, c=np.arange(101), cmap='Reds')\n",
    "ax.plot(x,y, marker = '.', ls = 'None', mfc = 'w', mec ='k')\n",
    "ax.quiver(x[:-1],y[:-1], x[1:]-x[:-1], y[1:]-y[:-1], np.arange(n), cmap = 'RdBu', \n",
    "          scale = 1, units = 'xy', scale_units = 'xy', width = 0.075, headwidth = 4)\n",
    "ax.set_xlim([-21,21])\n",
    "ax.set_ylim([-21,21])\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('y')\n",
    "ax.set_title('Radom walk on 2D square lattice')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate an ensemble of 2D walkers and coarse-grain to calculate entropy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "S = []\n",
    "ts = np.arange(10,5000,100)\n",
    "# create coarse grained grid for calculating Probabilities\n",
    "bin_edge = np.arange(-100,110,10)\n",
    "\n",
    "# loop over different max times\n",
    "for t in ts:\n",
    "    xh, yh = ensemble_2d(N = 1000, t_max = t,wall =200)\n",
    "    # h is counts/bin/total/Area, multiply h by bin area to get a probability\n",
    "    h, x_edges, y_edges = np.histogram2d(xh,yh, bins = [bin_edge, bin_edge], normed = True)\n",
    "    \n",
    "    # calculate entropy, taking care of log(0) with a small number\n",
    "    S.append(-1*np.sum(100*h*np.log(100*h+1e-40)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "ax.plot(ts,S, lw = 2)\n",
    "ax.set_xlabel('time step')\n",
    "ax.set_ylabel('Entropy')\n",
    "ax.set_title(\"Entropy in 2D square lattice\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check that we have a proper probability\n",
    "np.sum(h)*100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}