animating a stem plot in matplotlib

Question

I'm trying to animate a stem plot in matplotlib and I can't find the necessary documentation to help me. I have a series of data files which each look like this:

1 0.345346
2 0.124325
3 0.534585

and I want plot each file as a separate frame.

According to this and this other tutorial, I should create a function which updates the data contained in each plot object (artist? I'm not sure about the terminology)

From the second link, this is the update function

def update(frame):
global P, C, S

# Every ring is made more transparent
C[:,3] = np.maximum(0, C[:,3] - 1.0/n)

# Each ring is made larger
S += (size_max - size_min) / n

# Reset ring specific ring (relative to frame number)
i = frame % 50
P[i] = np.random.uniform(0,1,2)
S[i] = size_min
C[i,3] = 1

# Update scatter object
scat.set_edgecolors(C)
scat.set_sizes(S)
scat.set_offsets(P)

# Return the modified object
return scat,

How can I adapt this kind of update function for a stem plot? The documentation for stem is horribly brief (in fact this is a recurring issue as I'm learning matplotlib), but the example code shows that the output of stem is a tuple markerline, stemlines, baseline rather than an artist object like for plt.plot or plt.imshow.

So when I write my update function for the animation, how can I update the data inside the stem plot?

score 4 · Accepted Answer · answered Jul 06 '17 at 02:30

Here you go!

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

import numpy as np

fig, ax = plt.subplots()
x = np.linspace(0.1, 2*np.pi, 10)
markerline, stemlines, baseline = ax.stem(x, np.cos(x), '-.')

def update(i):
    ax.cla()
    markerline, stemlines, baseline = ax.stem(x, np.cos(x+i/10), '-.')
    ax.set_ylim((-1, 1))

anim = FuncAnimation(fig, update, frames=range(10, 110, 10), interval=500)
anim.save('so.gif', dpi=80, writer='imagemagick')

I think there can be better ways of achieving this- not requiring to clear the plot each time. However, this works!

Ok so you suggest to clear the plot and redraw at each frame. That will work, thanks! I'm finding it hard to learn matplotlib because the documentation is sparse. I can usually piece together the examples, but I don't really understand all of the mechanisms. For example, the difference between axes and artists, etc. Do you know a resource which explains the logic and structure in matplotlib? — Mr.Weathers, Jul 06 '17 at 17:43

scleronomic · Answer 2 · 2020-03-06T20:25:49.143

When using the keyword use_line_collection=True (default behavior since Matplotlib 3.3) one can update the three elements

markerline
stemlines
baseline

individualy. Here is the code for the sine wave example:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

fig, ax = plt.subplots()
x = np.linspace(0.1, 2*np.pi, 10)
y = np.cos(x)
bottom = 0
h_stem = ax.stem(x, y, bottom=bottom, use_line_collection=True, linefmt='-.')

def update(i):
    y = np.cos(x+i/10)

    # markerline
    h_stem[0].set_ydata(y)
    h_stem[0].set_xdata(x)  # not necessary for constant x 

    # stemlines
    h_stem[1].set_paths([np.array([[xx, bottom], 
                                   [xx, yy]]) for (xx, yy) in zip(x, y)])

    # baseline
    h_stem[2].set_xdata([np.min(x), np.max(x)])
    h_stem[2].set_ydata([bottom, bottom])  # not necessary for constant bottom

anim = FuncAnimation(fig, update, frames=range(10, 110, 10), interval=1)
anim.save('so.gif', dpi=80, writer='imagemagick')

Depending on what values (x, y, bottom) should be updated you can omit some parts of this update or reuse the current values. I wrote a more general function, where you can pass an arbitrary combination of these values:

def update_stem(h_stem, x=None, y=None, bottom=None):
    if x is None:
        x = h_stem[0].get_xdata()
    else:
        h_stem[0].set_xdata(x)
        h_stem[2].set_xdata([np.min(x), np.max(x)])

    if y is None:
        y = h_stem[0].get_ydata()
    else:
        h_stem[0].set_ydata(y)

    if bottom is None:
        bottom = h_stem[2].get_ydata()[0]
    else:
        h_stem[2].set_ydata([bottom, bottom])

    h_stem[1].set_paths([np.array([[xx, bottom], 
                                   [xx, yy]]) for (xx, yy) in zip(x, y)])

animating a stem plot in matplotlib

2 Answers2

Linked