Use interact widget in stem plot

Question

I am trying to make a Python stem plot in a Jupyter notebook.

The plot is this:

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
from scipy import fftpack

f = 10  
f_s = 100  
t = np.linspace(0, 2, 2 * f_s, endpoint=False)
x = np.sin(f * 2 * np.pi * t)

X = fftpack.fft(x)
freqs = fftpack.fftfreq(len(x)) * f_s 

fig, ax = plt.subplots()
ax.stem(freqs, np.abs(X))
ax.set_xlabel('Frecuency [Hz]')
ax.set_ylabel('Spectrum')
ax.set_xlim(-f_s / 2, f_s / 2)
ax.set_ylim(-5, 110);

And I would like to add a slider bar like in this post, Jupyter Notebook: interactive plot with widgets . The solution in this post

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
from scipy import fftpack
from ipywidgets import *

f = 10  
f_s = 100  
t = np.linspace(0, 2, 2 * f_s, endpoint=False)
x = np.sin(f * 2 * np.pi * t)

X = fftpack.fft(x)
freqs = fftpack.fftfreq(len(x)) * f_s 
fig, ax = plt.subplots()
line, = ax.stem(freqs, np.abs(X))
ax.set_xlabel('Frecuency [Hz]')
ax.set_ylabel('Spectrum')
ax.set_xlim(-f_s / 2, f_s / 2)
ax.set_ylim(-5, 110);

def update(f_s=100):
    line.set_ydata(fftpack.fftfreq(len(x)) * f_s)
    ax.set_xlim(-f_s / 2, f_s / 2)
    fig.canvas.draw()

interact(update);

returns me an error:

     ValueError                                Traceback (most recent call last)
<ipython-input-18-fd1aaf0ca96d> in <module>()
      7 
      8 fig, ax = plt.subplots()
----> 9 line, = ax.stem(freqs, np.abs(X))
     10 ax.set_xlabel('Frecuency [Hz]')
     11 ax.set_ylabel('Spectrum')

ValueError: too many values to unpack (expected 1)

Answer 1

I think there is foremost a problem with what you want to show. If you change the frequency, you would probably want to recalculate the spectrum as well. So let's correct that first, using a line plot.

%matplotlib notebook
import matplotlib.pyplot as plt
import numpy as np
from scipy import fftpack
from ipywidgets import *

def func(f_s):
    t = np.linspace(0, 2, 2 * f_s, endpoint=False)
    x = np.sin(f * 2 * np.pi * t)
    X = fftpack.fft(x)
    freqs = fftpack.fftfreq(len(x)) * f_s 
    return freqs, np.abs(X)

f = 10
f_s = 100  

fig, ax = plt.subplots()
line, = ax.plot(*func(f_s), marker="o")
ax.set_xlabel('Frecuency [Hz]')
ax.set_ylabel('Spectrum')
ax.set_xlim(-f_s / 2, f_s / 2)
ax.set_ylim(-5, 110);

def update(f_s=100):
    if f_s > 0:
        x, y = func(f_s)
        line.set_data(x, y)
        ax.set_xlim(-f_s / 2, f_s / 2)
        ax.set_ylim(None, 1.1*y.max())
        fig.canvas.draw()

interact(update);

Now the problem with a stem plot is that it would create a lot of lines which need to be updated. So I would rather recommend using a LineCollection.

%matplotlib notebook
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
import numpy as np
from scipy import fftpack
from ipywidgets import *

def func(f_s):
    t = np.linspace(0, 2, 2 * f_s, endpoint=False)
    x = np.sin(f * 2 * np.pi * t)
    X = fftpack.fft(x)
    freqs = fftpack.fftfreq(len(x)) * f_s 
    return freqs, np.abs(X)

f = 10
f_s = 100  

fig, ax = plt.subplots()
x,y = func(f_s)
markers, = ax.plot(x,y, ls="none", marker="o")
baseline = ax.axhline(0, color="crimson")

verts=np.c_[x, np.zeros_like(x), x, y].reshape(len(x),2,2)
col = LineCollection(verts)
ax.add_collection(col)

ax.set_xlabel('Frecuency [Hz]')
ax.set_ylabel('Spectrum')
ax.set_xlim(-f_s / 2, f_s / 2)
ax.set_ylim(None, 110);

def update(f_s=100):
    if f_s > 0:
        x, y = func(f_s)
        markers.set_data(x, y)
        verts=np.c_[x, np.zeros_like(x), x, y].reshape(len(x),2,2)
        col.set_segments(verts)
        ax.set_xlim(-f_s / 2, f_s / 2)
        ax.set_ylim(None, 1.1*y.max())
        fig.canvas.draw_idle()

interact(update);