Matplotlib Linear Regression Animation in Jupyter Notebook
By: B. Chen
Tags: Animation, Linear Regression, Data Visualization, Matplotlib, Python

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 NaN

df.dropna(subset=['Horsepower', 'Miles_per_Gallon'], inplace=True) # Transform data

x = 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()

Image by Author

2. Interactive Plot in Jupyter Notebook

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 plt

from matplotlib.animation import FuncAnimation

from 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

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

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.

, 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()

Matplotlib Linear Regression Animation in Jupyter Notebook (Image by Author)

Please check out Notebook for source code

3. Embedded HTML5 video in Jupyter Notebook

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 plt

from matplotlib.animation import FuncAnimation

from IPython import display # Turn off matplotlib plot in Notebook

plt.ioff()

# Pass the ffmpeg path

plt.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

Matplotlib Linear Regression Animation in Jupyter Notebook (Image by Author)

Please check out Notebook for source code

Conclusion

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.

You may be interested in some of my other Data Visualization and Pandas articles:

More tutorials can be found on my Github

References