{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Modeling gravitational fields"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.1 Defining a vector field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(x,y):\n",
    "    return (-2*y, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_vector_field(f,xmin,xmax,ymin,ymax,xstep=1,ystep=1):\n",
    "\n",
    "    X,Y = np.meshgrid(np.arange(xmin,xmax, xstep),np.arange(ymin,ymax, ystep))\n",
    "    U = np.vectorize(lambda x,y : f(x,y)[0])(X,Y)\n",
    "    V = np.vectorize(lambda x,y : f(x,y)[1])(X,Y)\n",
    "    plt.quiver(X, Y, U, V,color='red')\n",
    "    fig = plt.gcf()\n",
    "    fig.set_size_inches(7,7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32472 characters]\n",
      "text/plain": [
       "<Figure size 504x504 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_vector_field(f,-5,5,-5,5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2 Defining a simple force field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31052 characters]\n",
      "text/plain": [
       "<Figure size 504x504 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def f(x,y):\n",
    "    return (-x,-y)\n",
    "\n",
    "plot_vector_field(f,-5,5,-5,5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Adding gravity to the asteroid game"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See asteroids.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**EXERCISE** Where do all of the vectors in the vector field $(-2-x,4-y)$ point?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32588 characters]\n",
      "text/plain": [
       "<Figure size 504x504 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def g(x,y):\n",
    "    return (-2-x,4-y)\n",
    "plot_vector_field(g,-5,5,-5,5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3 Introducing potential energy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.1 Defining a potential energy scalar field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u(x,y):\n",
    "    return 0.5 * (x**2 + y**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "from matplotlib.ticker import LinearLocator, FormatStrFormatter\n",
    "\n",
    "def plot_scalar_field(f,xmin,xmax,ymin,ymax,xstep=0.25,ystep=0.25,c=None,cmap=cm.coolwarm,alpha=1,antialiased=False):\n",
    "    fig = plt.figure()\n",
    "    fig.set_size_inches(7,7)\n",
    "    ax = fig.gca(projection='3d')\n",
    "\n",
    "    fv = np.vectorize(f)\n",
    "    \n",
    "    # Make data.\n",
    "    X = np.arange(xmin, xmax, xstep)\n",
    "    Y = np.arange(ymin, ymax, ystep)\n",
    "    X, Y = np.meshgrid(X, Y)\n",
    "    Z = fv(X,Y)\n",
    "\n",
    "    # Plot the surface.\n",
    "    surf = ax.plot_surface(X, Y, Z, cmap=cmap,color=c,alpha=alpha,\n",
    "                           linewidth=0, antialiased=antialiased)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 63632 characters]\n",
      "text/plain": [
       "<Figure size 504x504 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_scalar_field(u, -5, 5, -5, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.2 Plotting a scalar field as a heatmap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Figure size 504x504 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 25320 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def scalar_field_heatmap(f,xmin,xmax,ymin,ymax,xstep=0.1,ystep=0.1):\n",
    "    fig = plt.figure()\n",
    "    fig.set_size_inches(7,7)\n",
    "\n",
    "    fv = np.vectorize(f)\n",
    "\n",
    "    X = np.arange(xmin, xmax, xstep)\n",
    "    Y = np.arange(ymin, ymax, ystep)\n",
    "    X, Y = np.meshgrid(X, Y)\n",
    "    \n",
    "    # https://stackoverflow.com/a/54088910/1704140\n",
    "    z = fv(X,Y)\n",
    "    \n",
    "#     # x and y are bounds, so z should be the value *inside* those bounds.\n",
    "#     # Therefore, remove the last value from the z array.\n",
    "#     z = z[:-1, :-1]\n",
    "#     z_min, z_max = -z.min(), z.max()\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    c = ax.pcolormesh(X,Y,z, cmap='plasma')\n",
    "    # set the limits of the plot to the limits of the data\n",
    "    ax.axis([X.min(), X.max(), Y.min(), Y.max()])\n",
    "    fig.colorbar(c, ax=ax)\n",
    "    plt.xlabel('x')\n",
    "    plt.ylabel('y')\n",
    "\n",
    "scalar_field_heatmap(u,-5,5,-5,5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.3 Plotting a scalar field as a contour map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 88328 characters]\n",
      "text/plain": [
       "<Figure size 504x504 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def scalar_field_contour(f,xmin,xmax,ymin,ymax,levels=None):\n",
    "\n",
    "    fv = np.vectorize(f)\n",
    "\n",
    "    X = np.arange(xmin, xmax, 0.1)\n",
    "    Y = np.arange(ymin, ymax, 0.1)\n",
    "    X, Y = np.meshgrid(X, Y)\n",
    "    \n",
    "    # https://stackoverflow.com/a/54088910/1704140\n",
    "    Z = fv(X,Y)\n",
    "    \n",
    "    fig, ax = plt.subplots()\n",
    "    CS = ax.contour(X, Y, Z,levels=levels)\n",
    "    ax.clabel(CS, inline=1, fontsize=10,fmt='%1.1f')\n",
    "    plt.xlabel('x')\n",
    "    plt.ylabel('y')\n",
    "    fig.set_size_inches(7,7)\n",
    "    \n",
    "scalar_field_contour(u,-10,10,-10,10,levels=[10,20,30,40,50,60])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4 Connecting energy and forces with the gradient"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.5 Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Plot the cross-section of $h(x,y) = e^y\\sin(x)$ where $y=1$.  Plot the cross section of $h(x,y)$ where $x=\\pi/6$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_function(f,xmin,xmax,color=None):\n",
    "    ts = np.linspace(xmin,xmax,1000)\n",
    "    plt.plot(ts,[f(t) for t in ts],c=color)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13792 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "from math import pi, e, sin\n",
    "plot_function(lambda x: e*sin(x),-pi,pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 9936 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import math\n",
    "plot_function(lambda y: math.exp(y)/2,-3,3)"
   ]
  },
  {
   "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
