matplotlib-linear-regression-animation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Matplotlib Linear Regression Animation in Jupyter Notebook\n",
    "\n",
    "This is a notebook for the medium article [Matplotlib Linear Regression Animation in Jupyter Notebook](https://bindichen.medium.com/matplotlib-linear-regression-animation-in-jupyter-notebook-2435b711bea2)\n",
    "\n",
    "Please check out article for instructions\n",
    "\n",
    "**License**: [BSD 2-Clause](https://opensource.org/licenses/BSD-2-Clause)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install vega_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'vega_datasets'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-1-3ee911dbcafe>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      2\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mnumpy\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[1;32mfrom\u001b[0m \u001b[0msklearn\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mlinear_model\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mLinearRegression\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 4\u001b[1;33m \u001b[1;32mfrom\u001b[0m \u001b[0mvega_datasets\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mdata\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      5\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      6\u001b[0m \u001b[0mget_ipython\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrun_line_magic\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'matplotlib'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'inline'\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'vega_datasets'"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from sklearn.linear_model import LinearRegression\n",
    "from vega_datasets import data\n",
    "\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'svg'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = data.cars()\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Drop rows with NaN\n",
    "df.dropna(subset=['Horsepower', 'Miles_per_Gallon'], inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Transform data\n",
    "x = df['Horsepower'].to_numpy().reshape(-1, 1)\n",
    "y = df['Miles_per_Gallon'].to_numpy().reshape(-1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(x, y, c='g', label='Horsepower vs. Miles_per_Gallon')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Animations with Interactive Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Enable interactive plot\n",
    "%matplotlib notebook\n",
    "%matplotlib notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.animation import FuncAnimation\n",
    "from sklearn.linear_model import LinearRegression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_data = []\n",
    "y_data = []\n",
    "\n",
    "fig, ax = plt.subplots()     # A tuple unpacking to unpack the only fig\n",
    "ax.set_xlim(30, 250)\n",
    "ax.set_ylim(5, 50)\n",
    "# Plotting \n",
    "scatter, = ax.plot([], [], 'go', label='Horsepower vs. Miles_per_Gallon')\n",
    "line, = ax.plot([], [], 'r', label='Linear Regression')\n",
    "ax.legend()\n",
    "\n",
    "reg = LinearRegression()\n",
    "\n",
    "def animate(frame_num):\n",
    "    # Adding data\n",
    "    x_data.append(x[frame_num])\n",
    "    y_data.append(y[frame_num])\n",
    "    # Convert data to numpy array\n",
    "    x_train = np.array(x_data).reshape(-1, 1)\n",
    "    y_train = np.array(y_data).reshape(-1, 1)\n",
    "    # Fit values to a linear regression\n",
    "    reg.fit(x_train, y_train)\n",
    "\n",
    "    # update data for scatter plot\n",
    "    scatter.set_data((x_data, y_data))\n",
    "    # Predict value and update data for line plot\n",
    "    line.set_data((list(range(250)), reg.predict(np.array([entry for entry in range(250)]).reshape(-1, 1))))\n",
    "\n",
    "anim = FuncAnimation(fig, animate, frames=len(x), interval=20)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Animations with embedded HTML5 video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.animation import FuncAnimation\n",
    "from IPython import display\n",
    "\n",
    "# Turn off matplotlib plot in Notebook\n",
    "plt.ioff()\n",
    "# Pass the ffmpeg path\n",
    "plt.rcParams['animation.ffmpeg_path'] = '/Users/chen5/Downloads/ffmpeg'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_data = []\n",
    "y_data = []\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_xlim(30, 250)\n",
    "ax.set_ylim(5, 50)\n",
    "scatter, = ax.plot([], [], 'go', label='Horsepower vs. Miles_per_Gallon')\n",
    "line, = ax.plot([], [], 'r', label='Linear Regression')\n",
    "ax.legend()\n",
    "\n",
    "reg = LinearRegression()\n",
    "\n",
    "def animate(frame_num):\n",
    "    # Adding data\n",
    "    x_data.append(x[frame_num])\n",
    "    y_data.append(y[frame_num])\n",
    "    # Convert data to numpy array\n",
    "    x_train = np.array(x_data).reshape(-1, 1)\n",
    "    y_train = np.array(y_data).reshape(-1, 1)\n",
    "    reg.fit(x_train, y_train)\n",
    "    \n",
    "    # update data for scatter plot\n",
    "    scatter.set_data((x_data, y_data))\n",
    "    # Predict value and update data for line plot\n",
    "    line.set_data((list(range(250)), reg.predict(np.array([entry for entry in range(250)]).reshape(-1, 1))))\n",
    "\n",
    "anim = FuncAnimation(fig, animate, frames=len(x), interval=20)\n",
    "\n",
    "video = anim.to_html5_video()\n",
    "html = display.HTML(video)\n",
    "display.display(html)\n",
    "plt.close()\n",
    "\n",
    "# Note Github only render static HTML and the embeded HTML5 video won't be displayed, \n",
    "# The embedded video should be working if you host the Notebook or open it locally. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Thanks for reading\n",
    "\n",
    "This is a notebook for the medium article [Matplotlib Linear Regression Animation in Jupyter Notebook](https://bindichen.medium.com/matplotlib-linear-regression-animation-in-jupyter-notebook-2435b711bea2)\n",
    "\n",
    "**License**: [BSD 2-Clause](https://opensource.org/licenses/BSD-2-Clause)"
   ]
  },
  {
   "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": 4
}