Cleaned up notebook.

Author: rbanerjee20

fdm-devito-notebooks/03_diffu/diffu_rw.ipynb

@@ -236,13 +236,40 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "This implementation is scalar, which is slow for large values of N. Thankfully, Devito gives us a simple way of implementing this using vectorization. Read on!"
+    "A random walk with $N$ steps, starting at $x_0$, where we move\n",
+    "to the left with probability $p$ and to the right\n",
+    "with probability $1-p$ can now be implemented by"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 303,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import random, numpy as np\n",
+    "\n",
+    "def random_walk1D(x0, N, p, random=random):\n",
+    "    \"\"\"1D random walk with 1 particle.\"\"\"\n",
+    "    # Store position in step k in position[k]\n",
+    "    position = np.zeros(N)\n",
+    "    position[0] = x0\n",
+    "    current_pos = x0\n",
+    "    for k in range(N-1):\n",
+    "        r = random.uniform(0, 1)\n",
+    "        if r <= p:\n",
+    "            current_pos -= 1\n",
+    "        else:\n",
+    "            current_pos += 1\n",
+    "        position[k+1] = current_pos\n",
+    "    return position"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "\n",
     "Since $N$ is supposed to be large and we want to repeat the process for\n",
     "many particles, we should speed up the code as much as possible.\n",
     "Vectorization is the obvious technique here: we draw all the random\n",
@@ -260,12 +287,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 295,
    "metadata": {},
    "outputs": [],
    "source": [
-    "import random, numpy as np\n",
-    "\n",
     "N=200\n",
     "p=50\n",
     "r = np.random.uniform(0, 1, size=N)\n",
@@ -276,14 +301,12 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "A random walk with $N$ steps, starting at $x_0$, where we move\n",
-    "to the left with probability $p$ and to the right\n",
-    "with probability $1-p$ can now be implemented by"
+    "Our vectorized implementation using Devito is as follows:"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 296,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -292,13 +315,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 64,
+   "execution_count": 304,
    "metadata": {},
    "outputs": [],
    "source": [
-    "def random_walk1D(x0, N, p, random=random):\n",
+    "def random_walk1D_vec(x0, N, p, random=random):\n",
     "    \"\"\"1D random walk with 1 particle.\"\"\"\n",
-    "    # Create and initialise r\n",
+    "    # Create and initialise r with dimension x_d\n",
     "    x_d = Dimension(name='x_d', spacing=Constant('h_x'))\n",
     "    r = TimeFunction(name='r', dimensions=(x_d,), shape=(N+1,))\n",
     "    r.data[0] = x0\n",
@@ -329,7 +352,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 259,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -648,7 +671,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 216,
+   "execution_count": 283,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -731,7 +754,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 39,
+   "execution_count": 297,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -740,14 +763,14 @@
     "    # produce the same result\n",
     "    x0 = 0;  N = 4;  p = 0.5\n",
     "\n",
-    "    # First, check that random_walks1D for 1 walk reproduces\n",
+    "    # Check that random_walks1D for 1 walk reproduces\n",
     "    # the walk in random_walk1D\n",
     "    num_walks = 1\n",
     "    np.random.seed(10)\n",
     "    computed = random_walks1D(\n",
     "        x0, N, p, num_walks, random=np.random)\n",
     "    np.random.seed(10)\n",
-    "    expected = random_walk1D(\n",
+    "    expected = random_walk1D_vec(\n",
     "        x0, N, p, random=np.random)\n",
     "    print(\"Computed: \", computed[0])\n",
     "    print(\"Expected: \", expected)\n",
@@ -756,7 +779,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 40,
+   "execution_count": 298,
    "metadata": {},
    "outputs": [
     {
@@ -857,7 +880,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 43,
+   "execution_count": 239,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1061,7 +1084,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 211,
+   "execution_count": 299,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1121,7 +1144,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 212,
+   "execution_count": 301,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1152,7 +1175,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 219,
+   "execution_count": 305,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1163,11 +1186,11 @@
     "    x_d = Dimension(name='x_d', spacing=Constant('h_x'))\n",
     "    # Add space dimensions\n",
     "    d_d = Dimension(name='d', spacing=Constant('h_d'))\n",
-    "    r = TimeFunction(name='r', dimensions=(x_d,d_d), shape=(N+1,d))\n",
+    "    r = TimeFunction(name='r', dimensions=(x_d, d_d), shape=(N+1, d))\n",
     "    r.data[0] = x0\n",
     "    \n",
     "    rs = random.uniform(0, 1, size=N*d).reshape(N,d)\n",
-    "    steps = Function(name='steps', dimensions=(x_d,d_d), shape=(N+1,d))\n",
+    "    steps = Function(name='steps', dimensions=(x_d, d_d), shape=(N+1, d))\n",
     "    steps.data[:N] = np.where(rs <= p, -1, 1)\n",
     "    \n",
     "    # The next value is computed using the previous value and the steps function at that point\n",
@@ -1198,7 +1221,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 218,
+   "execution_count": 225,
    "metadata": {},
    "outputs": [],
    "source": [