22

I have this code. I want to add a subplot to draw the cosine function. (I do not want to create a class). The second plot should be dynamically updated as well

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

def data_gen():
    t = data_gen.t
    cnt = 0
    while cnt < 1000:
        cnt+=1
        t += 0.05
        yield t, np.sin(2*np.pi*t) * np.exp(-t/10.)
data_gen.t = 0

fig, ax = plt.subplots()
line, = ax.plot([], [], lw=2)
ax.set_ylim(-1.1, 1.1)
ax.set_xlim(0, 5)
ax.grid()
xdata, ydata = [], []
def run(data):
    # update the data
    t,y = data
    xdata.append(t)
    ydata.append(y)
    xmin, xmax = ax.get_xlim()

    if t >= xmax:
        ax.set_xlim(xmin, 2*xmax)
        ax.figure.canvas.draw()
    line.set_data(xdata, ydata)

    return line,

ani = animation.FuncAnimation(fig, run, data_gen, blit=True, interval=10,
    repeat=False)
plt.show()
omar_9
  • 19
  • 1
  • 6
mv93
  • 265
  • 1
  • 4
  • 13

1 Answers1

42

Basically you can use a very similar structure as the one you have in your example. You only need to create an additional axes (subplot) and a second line object:

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

def data_gen():
    t = data_gen.t
    cnt = 0
    while cnt < 1000:
        cnt+=1
        t += 0.05
        y1 = np.sin(2*np.pi*t) * np.exp(-t/10.)
        y2 = np.cos(2*np.pi*t) * np.exp(-t/10.)
        # adapted the data generator to yield both sin and cos
        yield t, y1, y2

data_gen.t = 0

# create a figure with two subplots
fig, (ax1, ax2) = plt.subplots(2,1)

# intialize two line objects (one in each axes)
line1, = ax1.plot([], [], lw=2)
line2, = ax2.plot([], [], lw=2, color='r')
line = [line1, line2]

# the same axes initalizations as before (just now we do it for both of them)
for ax in [ax1, ax2]:
    ax.set_ylim(-1.1, 1.1)
    ax.set_xlim(0, 5)
    ax.grid()

# initialize the data arrays 
xdata, y1data, y2data = [], [], []
def run(data):
    # update the data
    t, y1, y2 = data
    xdata.append(t)
    y1data.append(y1)
    y2data.append(y2)

    # axis limits checking. Same as before, just for both axes
    for ax in [ax1, ax2]:
        xmin, xmax = ax.get_xlim()
        if t >= xmax:
            ax.set_xlim(xmin, 2*xmax)
            ax.figure.canvas.draw()

    # update the data of both line objects
    line[0].set_data(xdata, y1data)
    line[1].set_data(xdata, y2data)

    return line

ani = animation.FuncAnimation(fig, run, data_gen, blit=True, interval=10,
    repeat=False)
plt.show()

enter image description here

Trenton McKinney
  • 56,955
  • 33
  • 144
  • 158
hitzg
  • 12,133
  • 52
  • 54
  • could you explain the difference between line1, = ax1.plot([], [], lw=2) and line1 = ax1.plot([], [], lw=2) (second version missing comma)? – ZakS May 10 '18 at 11:58
  • 3
    `ax1.plot` returns a list of (potentially) multiple lines. With the comma, the first element of that list is assigned to `line1`. Without the comma `line1` will contain the entire list. – hitzg May 12 '18 at 14:38
  • How to modify this one if y1 and y2 need to be drawn on the same subplot but the data_gen is a function with some input parameter k. So each subplot show for different values of k. – Yograj Singh Mandloi Apr 26 '19 at 10:18
  • How to run this in jupyter notebook? adding `%matplotlib inline` doesn't work – MattS Oct 08 '19 at 12:38
  • @MattS `%matplotlib notebook` – Indivara Oct 14 '19 at 14:15
  • with the line "line1, = ax1.plot([], [], lw=2)" I've got an error of too many values to unpack (expected 1)... if I remove the comma at line "line[0].set_data(xdata, y1data) " I've got error 'tuple' object has no attribute 'set_data' – Jonathan Roy Feb 11 '22 at 20:40