I'm working on a fairly simple matplotlib animation using animation.FuncAnimation and it takes a very long time to run; about 0.81s per frame and 180s for the 220 frames in an 11 second long video. I've already incorporated the standard performance improvements, namely not creating Artist objects in the animation loop and returning all updated Artists objects so that bliting can be done but neither has had a significant impact on performance. I've timed the contents of the newFrame function and it only takes about 0.6ms to run. My code is below.
Any other suggestions how how to speed this up or parallelize it?
EDIT: It looks like the performance issues are largely to do with the subplots. Reducing the number of subplots reduces the execution time dramatically. Not really a great solution but noteworthy
EDIT 2: Commenting out axsArray[a,b].minorticks_on() doubles the performance to ~0.4s per frame. This brings performance more in line with what I expect based on previous experience with matplotlib animations, though still incredibly slow for a simple plot.
#!/usr/bin/env python3
from timeit import default_timer
import matplotlib
import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np
import os
import sys
matplotlib.use("Agg")
plt.close('all')
start = default_timer()
# ==============================================================================
def main():
# ==========================================================================
# Settings
# ==========================================================================
# Check for CLI arguments
if len(sys.argv) == 2:
cliFPS = int(sys.argv[1])
elif len(sys.argv) > 2:
raise TypeError(f"Too many ({len(sys.argv)}) command line arguments given")
# Data settings
loadPath = ''
yPos = 16 # y position of the x line to load
zPos = 16 # z position of the x line to load
# Plot Settings
supTitleText = "Time Evolution of Initial Conditions"
densityColor = 'blue' # color of the density plot
velocityColor = 'purple' # color of the velocity plots
magneticColor = 'tab:orange' # color of the magnetic field plots
pressureColor = 'green' # color of the pressure plot
ieColor = 'red' # color of the specific internal energy plot
linestyle = '-' # The line style
linewidth = 0.5 # How wide to make the lines
marker = "." # Marker kind for points
markersize = 3 # Size of the marker
figSizeScale = 2. # Scaling factor for the figure size
figHeight = 4.8 * figSizeScale # height of the plot in inches, default is 4.8
figWidth = 7.0 * figSizeScale # width of the plot in inches, default is 6.4
padPercent = 0.05 # How many percent larger the limits should be than the data
# Video Settings
OutFile = "mvp.mp4" # Output filename
Duration = 10. # How long the video is in seconds
dpi = 150 # Dots per inch
index = 0 # Initialize index
initIndex = 0 # Index for init frames
fps = cliFPS if ("cliFPS" in locals()) else 20 # Framerate
FrameTime = (1./fps) * 1000 # Frametime in milliseconds
totFrames = int(fps * Duration) # Total number of frames (floor)
# ==========================================================================
# End settings
# ==========================================================================
# Load data
(densityData, velocityXData, velocityYData, velocityZData, pressureData,
ieData, magneticXData, magneticYData, magneticZData, positions,
timeStepNum, dims, physicalSize) = loadData(loadPath, yPos, zPos)
# Compute which time steps to plot
if timeStepNum.size >= totFrames:
floatSamples = np.arange(0, timeStepNum.size, timeStepNum.size/totFrames)
timeStepSamples = np.asarray(np.floor(floatSamples), dtype="int")
else: # if the number of simulation steps is less than the total number of frames
totFrames = timeStepNum.size
fps = np.ceil(totFrames/Duration)
FrameTime = (1./fps) * 1000
timeStepSamples = np.arange(0, timeStepNum.size, 1, dtype="int")
# Insert the initial second of the initial conditions
timeStepSamples = np.insert(timeStepSamples, 0, [0]*fps)
# Get the plot limits
densityLowLim, densityHighLim = computeLimit(densityData, padPercent)
ieLowLim, ieHighLim = computeLimit(ieData, padPercent)
pressureLowLim, pressureHighLim = computeLimit(pressureData, padPercent)
velocityXLowLim, velocityXHighLim = computeLimit(velocityXData, padPercent)
velocityYLowLim, velocityYHighLim = computeLimit(velocityYData, padPercent)
velocityZLowLim, velocityZHighLim = computeLimit(velocityXData, padPercent)
magneticXLowLim, magneticXHighLim = computeLimit(magneticXData, padPercent)
magneticYLowLim, magneticYHighLim = computeLimit(magneticYData, padPercent)
magneticZLowLim, magneticZHighLim = computeLimit(magneticZData, padPercent)
# Set up plots
# Create 9 subplots
fig, subPlot = plt.subplots(3, 3, figsize = (figWidth, figHeight))
# Super Title
titleText = fig.suptitle(supTitleText)
# Shared x-label
subPlot[2,0].set_xlabel("Position")
subPlot[2,1].set_xlabel("Position")
subPlot[2,2].set_xlabel("Position")
# Set values for the subplots
densityPlot, subPlot = setupSubPlots('Density', densityLowLim, densityHighLim, positions, densityData, linestyle, linewidth, marker, markersize, densityColor, subPlot)
pressurePlot, subPlot = setupSubPlots('Pressure', pressureLowLim, pressureHighLim, positions, pressureData, linestyle, linewidth, marker, markersize, pressureColor, subPlot)
iePlot, subPlot = setupSubPlots('Internal Energy', ieLowLim, ieHighLim, positions, ieData, linestyle, linewidth, marker, markersize, ieColor, subPlot)
velocityXPlot, subPlot = setupSubPlots('$V_x$', velocityXLowLim, velocityXHighLim, positions, velocityXData, linestyle, linewidth, marker, markersize, velocityColor, subPlot)
velocityYPlot, subPlot = setupSubPlots('$V_y$', velocityYLowLim, velocityYHighLim, positions, velocityYData, linestyle, linewidth, marker, markersize, velocityColor, subPlot)
velocityZPlot, subPlot = setupSubPlots('$V_z$', velocityZLowLim, velocityZHighLim, positions, velocityZData, linestyle, linewidth, marker, markersize, velocityColor, subPlot)
magneticXPlot, subPlot = setupSubPlots('$B_x$', magneticXLowLim, magneticXHighLim, positions, magneticXData, linestyle, linewidth, marker, markersize, magneticColor, subPlot)
magneticYPlot, subPlot = setupSubPlots('$B_y$', magneticYLowLim, magneticZHighLim, positions, magneticYData, linestyle, linewidth, marker, markersize, magneticColor, subPlot)
magneticZPlot, subPlot = setupSubPlots('$B_z$', magneticZLowLim, magneticZHighLim, positions, magneticZData, linestyle, linewidth, marker, markersize, magneticColor, subPlot)
# Layout
plt.tight_layout()
fig.subplots_adjust(top=0.88)
# Generate animation
simulation = animation.FuncAnimation(fig,
newFrame,
fargs = (fig, supTitleText, densityPlot, pressurePlot, iePlot,
velocityXPlot, velocityYPlot, velocityZPlot,
magneticXPlot, magneticYPlot, magneticZPlot,
densityData, pressureData, ieData,
velocityXData, velocityYData, velocityZData,
magneticXData, magneticYData, magneticZData,
timeStepSamples, titleText),
blit = True,
frames = timeStepSamples,
interval = FrameTime,
repeat = False)
FFwriter = animation.FFMpegWriter(bitrate=1000,
fps=fps,
codec='libx264',
extra_args=['-crf','28','-preset','ultrafast','-pix_fmt','yuv420p'])
simulation.save(filename=OutFile, writer = FFwriter)
# simulation.save(filename=OutFile, fps=fps, dpi=dpi)
print(f"\n\nAnimation complete. Framerate: {fps} fps, Total Number of Frames: {timeStepSamples.size}")
# ==============================================================================
# ==============================================================================
def loadData(path, yPos, zPos):
numFiles = 870
dims = [32, 32, 32]
physicalSize = [1,1,1]
np.random.random((numFiles, dims[0]))
# Allocate Arrays
densityData = np.random.random((numFiles, dims[0]))
velocityXData = np.random.random((numFiles, dims[0]))
velocityYData = np.random.random((numFiles, dims[0]))
velocityZData = np.random.random((numFiles, dims[0]))
pressureData = np.random.random((numFiles, dims[0]))
ieData = np.random.random((numFiles, dims[0]))
magneticXData = np.random.random((numFiles, dims[0]))
magneticYData = np.random.random((numFiles, dims[0]))
magneticZData = np.random.random((numFiles, dims[0]))
timeStepNum = np.arange(0, numFiles, 1)
positions = np.linspace(0., 1, dims[0])
return (densityData, velocityXData, velocityYData, velocityZData,
pressureData, ieData, magneticXData, magneticYData, magneticZData,
positions, timeStepNum, dims, physicalSize)
# ==============================================================================
# ==============================================================================
def setupSubPlots(fieldName, lowLim, highLim, positions, data, linestyle, linewidth, marker, markersize, color, axsArray):
# Get the subplot coordinates to set
if fieldName == 'Density':
a, b = (0,0)
elif fieldName == 'Pressure':
a, b = (0,1)
elif fieldName == 'Internal Energy':
a, b = (0,2)
elif fieldName == '$V_x$':
a, b = (1,0)
elif fieldName == '$V_y$':
a, b = (1,1)
elif fieldName == '$V_z$':
a, b = (1,2)
elif fieldName == '$B_x$':
a, b = (2,0)
elif fieldName == '$B_y$':
a, b = (2,1)
elif fieldName == '$B_z$':
a, b = (2,2)
else:
raise ValueError('setSubPlots received invalid fieldName')
# Set plot parameters
axsArray[a,b].set_ylim(lowLim, highLim)
axsArray[a,b].set_ylabel(fieldName)
axsArray[a,b].minorticks_on()
axsArray[a,b].grid(which = "both")
# Set initial values
returnPlot, = axsArray[a,b].plot(positions,
data[0,:],
linestyle = linestyle,
linewidth = linewidth,
marker = marker,
markersize = markersize,
color = color,
label = fieldName,
animated = True)
return returnPlot, axsArray
# ==============================================================================
# ==============================================================================
def computeLimit(dataSet, padPercent):
pad = np.max(np.abs([dataSet.min(), dataSet.max()])) * padPercent
if pad == 0:
# if the dataset doesn't exist or is zero return reasonable limits
return -1, 1
else:
lowLim = dataSet.min() - pad
highLim = dataSet.max() + pad
return lowLim, highLim
# ==============================================================================
# ==============================================================================
def newFrame(idx, fig, supTitleText, densityPlot, pressurePlot, iePlot,
velocityXPlot, velocityYPlot, velocityZPlot,
magneticXPlot, magneticYPlot, magneticZPlot,
densityData, pressureData, ieData,
velocityXData, velocityYData, velocityZData,
magneticXData, magneticYData, magneticZData,
timeStepSamples, titleText):
titleText.set_text(f"{supTitleText} \n Time Step: {idx}")
densityPlot .set_ydata(densityData[idx,:])
pressurePlot .set_ydata(pressureData[idx,:])
iePlot .set_ydata(ieData[idx,:])
velocityXPlot.set_ydata(velocityXData[idx,:])
velocityYPlot.set_ydata(velocityYData[idx,:])
velocityZPlot.set_ydata(velocityZData[idx,:])
magneticXPlot.set_ydata(magneticXData[idx,:])
magneticYPlot.set_ydata(magneticYData[idx,:])
magneticZPlot.set_ydata(magneticZData[idx,:])
# Report progress
if not hasattr(newFrame, "counter"):
newFrame.counter = -1 # Accounts for first call which is performed before the animation starts
print()
if newFrame.counter >= 0:
print(f'Animation is {100*(newFrame.counter/timeStepSamples.shape[0]):.1f}% complete', end='\r')
newFrame.counter += 1
# The return is required to make blit work
return (titleText, densityPlot, pressurePlot, iePlot,
velocityXPlot, velocityYPlot, velocityZPlot,
magneticXPlot, magneticYPlot, magneticZPlot)
# ==============================================================================
main()
print(f'Time to execute: {round(default_timer()-start,2)} seconds')
