Matplotlib Linear Regression Animation in Jupyter Notebook | by B. Chen | Towards Data Science

When fitting values to a line using Linear Regression, it can be very helpful to illustrate how the line fits the data as more data are added. In this article, you’ll learn how to create a Matplotlib animation, this article extends the topic from the previous article “animating a simple sine wave in Jupyter Notebook” to fit values to a line using Linear Regression.
Matplotlib is one of the most popular plotting libraries for exploratory data analysis. Plotting a static graph should work well in most cases, but when you are running simulations or doing time-series data analysis, basic plots may not always be enough. You may want to show an animation that helps you understand how the state changes over time. In this article, you’ll learn how to create a Matplotlib animation for fitting values to a line using Linear Regression. This article is structured as follows:
Gathering and preparing dataInteractive Plot in Jupyter NotebookEmbedded HTML5 video in Jupyter Notebook
Gathering and preparing data
Interactive Plot in Jupyter Notebook
Embedded HTML5 video in Jupyter Notebook
Please check out the Notebook for source code. More tutorials are available from Github Repo.
For demonstration, we will be using the cars dataset. The dataset includes the names of the cars, mpg, cylinders, displacement, etc. (8 variables) for 406 different cars [1]. We will be loading the data from Vega Datasets. Let’s import some libraries and load data to get started:
import matplotlib.pyplot as pltimport numpy as npfrom sklearn.linear_model import LinearRegressionfrom vega_datasets import data%matplotlib inline%config InlineBackend.figure_format = 'svg'df = data.cars()df.head()
For this tutorial, we will be plotting Horsepower vs. Miles_per_Gallon. Let’s drop rows with NaN and transform the data into proper format:
# Drop rows with NaNdf.dropna(subset=['Horsepower', 'Miles_per_Gallon'], inplace=True)# Transform datax = df['Horsepower'].to_numpy().reshape(-1, 1)y = df['Miles_per_Gallon'].to_numpy().reshape(-1, 1)
Let’s do a simple scatter plot to show Horsepower vs. Miles_per_Gallon:
plt.scatter(x, y, c='g', label='Horsepower vs. Miles_per_Gallon')plt.legend()plt.show()
In order to create an interactive plot in Jupyter Notebook, you first need to enable interactive plot as follows:
# Enable interactive plot%matplotlib notebook
After that, we import the required libraries. Especially FuncAnimation class for creating animation and LinearRegression for creating a Linear Regression model.
import matplotlib.pyplot as pltfrom matplotlib.animation import FuncAnimationfrom sklearn.linear_model import LinearRegression
Next, we need to create the initial state of the animation figure. Both x_data and y_data are initialized to []. We call subplots() without any arguments to create a Figure fig and a single Axes ax. We set the x range to (30, 250) and y range to (5, 50) to avoid them to be constantly changing. ax.plot() is called twice to create a scatter plot and a line plot. We also call LinearRegression() to create a Linear Regression model reg.
x_data = []y_data = []fig, ax = plt.subplots() ax.set_xlim(30, 250)ax.set_ylim(5, 50)scatter, = ax.plot([], [], 'go', label='Horsepower vs. Miles_per_Gallon')line, = ax.plot([], [], 'r', label='Linear Regression')ax.legend()reg = LinearRegression()
We then create a function animate() that’s going to be called by the FuncAnimation(). The function takes one argument frame_num — the current frame number. What we want to do here is to change data for our scatter plot and linear regression plot according to the frame number. Notice that reg.fit() is called to fit values as more data is added.
def animate(frame_num):    # Adding data    x_data.append(x[frame_num])    y_data.append(y[frame_num])    # Convert data to numpy array    x_train = np.array(x_data).reshape(-1, 1)    y_train = np.array(y_data).reshape(-1, 1)    # Fit values to a linear regression    reg.fit(x_train, y_train)    # update data for scatter plot    scatter.set_data((x_data, y_data))    # Predict value and update data for line plot    line.set_data((list(range(250)), reg.predict(np.array([entry for entry in range(250)]).reshape(-1, 1))))
Finally, we create our animation object by calling FuncAnimation with 4 arguments
The first argument fig is the reference to the figure we created
The second argument animate is the function we created to call at each frame to update the plot
The third argument frames=len(x), it defines the number of frames for “one round of animation” and we have set it to the number of training data.
Finally, the interval=20 argument sets the delay (in milliseconds) between frames. 20 is equivalent to 50FPS (1000ms / 20 = 50 FPS). If the number is too large you wait a really long time, if the number is too small, it would be faster than your eyes could see. In general, we need FPS greater than 16 for smooth animation (Human eyes can only receive 10–12 frames [2]).
anim = FuncAnimation(fig, animate, frames=len(x), interval=20)plt.show()
Please check out Notebook for source code
In the previous example, we have created a nice animation with linear regression. However, the plot is animating only when the code is running. Of course, we could take a screen capture, but that’s not efficient when you want to share your Jupyter Notebook online.
What we can do is convert the animation into an HTML5 video and embed it in Jupyter Notebook. We will be using FFmpeg for conversion. If you don’t have it, you first need to follow the instruction to download FFmpeg and unzip it.
After that, we import the required libraries and set 'ffmpeg_path' to the path to your local ffmpeg executable:
import matplotlib.pyplot as pltfrom matplotlib.animation import FuncAnimationfrom IPython import display# Turn off matplotlib plot in Notebookplt.ioff()# Pass the ffmpeg pathplt.rcParams['animation.ffmpeg_path'] = '/path_to_your/ffmpeg'
Creating animation is the same as the previous example.
x_data = []y_data = []fig, ax = plt.subplots()ax.set_xlim(30, 250)ax.set_ylim(5, 50)scatter, = ax.plot([], [], 'go', label='Horsepower vs. Miles_per_Gallon')line, = ax.plot([], [], 'r', label='Linear Regression')ax.legend()reg = LinearRegression()def animate(frame_num):    # Adding data    x_data.append(x[frame_num])    y_data.append(y[frame_num])    # Convert data to numpy array    x_train = np.array(x_data).reshape(-1, 1)    y_train = np.array(y_data).reshape(-1, 1)    reg.fit(x_train, y_train)        # update data for scatter plot    scatter.set_data((x_data, y_data))    # Predict value and update data for line plot    line.set_data((list(range(250)), reg.predict(np.array([entry for entry in range(250)]).reshape(-1, 1))))anim = FuncAnimation(fig, animate, frames=len(x), interval=20)
But instead of plt.show() to plot it, we are calling anim.to_html5_video() method to convert the animation result into an HTML5 video. We then need to get the HTML code that does the embedding for this video and that is done by calling IPython display.HTML(video). Finally, we call display.display(html) to embed the HTML code in Jupyter Notebook.
video = anim.to_html5_video()html = display.HTML(video)display.display(html)plt.close()               # avoid plotting a spare static plot
Please check out Notebook for source code
In this article, we have learned 2 approaches to create a Matplotlib Linear Regression animation in Jupyter Notebook. Creating an animation plot can help you running simulations and doing time-series data analysis.
I hope this article will help you to save time in learning Matplotlib. I recommend you to check out the documentation for more options & settings and to know about other things you can do.
Thanks for reading. Please check out the notebook for the source code and stay tuned if you are interested in the practical aspect of machine learning.
Matplotlib animations in Jupyter Notebook
Python Interactive Data Visualization with Altair
Interactive Data Visualization for exploring Coronavirus Spreads
10 tricks for converting numbers and strings to Datetime in Pandas
All Pandas json_normalize() you should know for flattening JSON
Using Pandas method chaining to improve code readability
How to do a Custom Sort on Pandas DataFrame
All the Pandas shift() you should know for data analysis
More tutorials can be found on my Github
[1] Dataset — cars http://lib.stat.cmu.edu/datasets/