{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 06a: Using Matplotlib to Create Animation\n",
    "\n",
    "This notebooks shows how to create an animation that will run and play inside a notebook.  The animation can also be written to a video file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import matplotlib.animation as animation\n",
    "from IPython.display import HTML\n",
    "\n",
    "import cartopy\n",
    "import cartopy.crs as ccrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_y(x, a, b, c, d):\n",
    "    y = a * np.sin(b * x + c) + d\n",
    "    return y\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "x = np.linspace(0, 2 * np.pi, 100)\n",
    "\n",
    "a = 1\n",
    "b = 1\n",
    "c = 0\n",
    "d = 0\n",
    "y0 = get_y(x, a, b, c, d)\n",
    "\n",
    "line = ax.plot(x, y0, 'b-', clip_on=False)\n",
    "line = line[0]\n",
    "\n",
    "xd = np.pi\n",
    "yd = get_y(xd, a, b, c, d)\n",
    "dot = ax.plot(xd, yd, 'ro', ms=10, clip_on=False)\n",
    "dot = dot[0]\n",
    "\n",
    "ax.set_xlim(0, 2 * np.pi)\n",
    "ax.set_ylim(-a, a)\n",
    "\n",
    "frames = 50\n",
    "\n",
    "def update(i):\n",
    "\n",
    "    dt = 2 * np.pi / frames\n",
    "    c = i * dt\n",
    "\n",
    "    # for each frame, update the data stored on each artist.\n",
    "    y = get_y(x, a, b, c, d)\n",
    "    line.set_xdata(x)\n",
    "    line.set_ydata(y)\n",
    "\n",
    "    yd = get_y(xd, a, b, c, d)\n",
    "    dot.set_ydata([yd])\n",
    "\n",
    "    return line,\n",
    "\n",
    "ani = animation.FuncAnimation(fig=fig, func=update, frames=frames)\n",
    "plt.close()\n",
    "\n",
    "\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the following can be used to save the animation to a file\n",
    "ext = \"mp4\" # mac mp4\n",
    "# ext = \"avi\" # windows avi\n",
    "ani.save(f\"sin_animation.{ext}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## a map example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting data on maps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating some data that spans the whole globe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lon, lat = np.mgrid[-180:181, -90:91]\n",
    "\n",
    "# build a fake dataset that shifts through n time steps\n",
    "n = 45\n",
    "data = np.array([2 * np.sin(3 * np.deg2rad(lon+dt)) + 3 * np.cos(4 * np.deg2rad(lat+dt))\n",
    "                 for dt in np.arange(0,n)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's check the first time step\n",
    "plt.pcolormesh(lon, lat, data[0])\n",
    "plt.colorbar();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's animate it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "mesh = ax.pcolormesh(lon, lat, data[0])\n",
    "fig.colorbar(mesh, orientation='vertical')\n",
    "\n",
    "def animate(i):\n",
    "    h = data[i]\n",
    "    mesh.set_array(h)\n",
    "    title = ax.set_title(f'Year {i+1}')\n",
    "    return mesh\n",
    "\n",
    "ani = animation.FuncAnimation(\n",
    "    fig, \n",
    "    animate, \n",
    "    frames=data.shape[0],\n",
    "    # save_count=data.shape[0],\n",
    "    # interval=500,\n",
    ")\n",
    "plt.close()\n",
    "\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Geospatially"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plotting data on a Cartesian grid is equivalent to plotting data in the PlateCarree projection, where meridians and parallels are all straight lines with constant spacing. As a result of this simplicity, global datasets we use often begin in the PlateCarree projection.\n",
    "\n",
    "We can plot these data values as a contour on a `GeoAxes` map. To do so, we must specify the `transform` keyword argument, which specifies the projection type used by our data. The `transform` keyword can be given to all matplotlib plotting methods. The projection type specified in `pcolormesh` will be transformed into the projection type specified in the `subplot` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(11, 8.5))\n",
    "ax = plt.subplot(1, 1, 1, projection=ccrs.Orthographic(central_latitude=45))\n",
    "ax.coastlines()\n",
    "mesh = ax.pcolormesh(lon, lat, data[0], transform=ccrs.PlateCarree())\n",
    "plt.colorbar(mesh, orientation='horizontal', pad=.05);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Class Activity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Animate the above `GeoAxes` in the way we had done with the simpler graph previously. \n",
    "\n",
    "Bonus: How does it work if you switch over to using `contourf` instead?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pyclass",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}