matplotlib 茎图的优化

Question

我正在尝试使用“matplotlib.pyplot.stem”函数生成一个 Stem 图。该代码有效，但处理时间超过 5 分钟。

我在 Matlab 中有一个类似的代码，几乎可以立即用相同的输入数据生成相同的绘图。

有没有办法优化此代码以提高速度或我可以使用的更好的功能？

茎图“H”和“plotdata”的参数是 16384 x 1 数组。

def stemplot():

    import numpy as np
    from scipy.fftpack import fft
    import matplotlib.pyplot as plt

    ################################################
    # Code to set up the plot data

    N=2048
    dr = 100

    k = np.arange(0,N)

    cos = np.cos
    pi = np.pi

    w = 1-1.932617*cos(2*pi*k/(N-1))+1.286133*cos(4*pi*k/(N-1))-0.387695*cos(6*pi*k/(N-1))+0.0322227*cos(8*pi*k/(N-1))

    y = np.concatenate([w, np.zeros((7*N))])

    H = abs(fft(y, axis = 0))
    H = np.fft.fftshift(H)
    H = H/max(H)
    H = 20*np.log10(H)
    H = dr+H 
    H[H < 0] = 0        # Set all negative values in dr+H to 0

    plotdata = ((np.arange(1,(8*N)+1,1))-1-4*N)/8
    #################################################

    # Plotting Code

    plt.figure
    plt.stem(plotdata,H,markerfmt = " ")

    plt.axis([(-4*N)/8, (4*N)/8, 0, dr])    
    plt.grid()
    plt.ylabel('decibels')
    plt.xlabel('DFT bins')
    plt.title('Frequency response (Flat top)')
    plt.show()


    return

这里也是matlab代码供参考：

N=2048;
dr = 100;
k=0:N-1

w = 1 - 1.932617*cos(2*pi*k/(N-1)) + 1.286133*cos(4*pi*k/(N-1)) -0.387695*cos(6*pi*k/(N-1)) +0.0322227*cos(8*pi*k/(N-1));

H = abs(fft([w zeros(1,7*N)]));
H = fftshift(H);
H = H/max(H);
H = 20*log10(H);
H = max(0,dr+H); % Sets negative numbers in dr+H to 0


figure
stem(([1:(8*N)]-1-4*N)/8,H,'-');
set(findobj('Type','line'),'Marker','none','Color',[.871 .49 0])
xlim([-4*N 4*N]/8)
ylim([0 dr])
set(gca,'YTickLabel','-100|-90|-80|-70|-60|-50|-40|-30|-20|-10|0')
grid on
ylabel('decibels')
xlabel('DFT bins')
title('Frequency response (Flat top)')

Answer 1

您可以使用ax.vlines 以您想要的格式模拟一个茎图。写一个小函数，

def make_stem(ax, x, y, **kwargs):
    ax.axhline(x[0],x[-1],0, color='r')

    ax.vlines(x, 0, y, color='b')

    ax.set_ylim([1.05*y.min(), 1.05*y.max()])

然后将示例中的相关行更改如下：

    # Plotting Code

##    plt.figure
##    plt.stem(plotdata,H,markerfmt = " ")

##    plt.axis([(-4*N)/8, (4*N)/8, 0, dr])

    fig, ax = plt.subplots()
    make_stem(ax, plotdata, H)

或多或少会立即产生情节。但是，我不知道这是否比@ImportanceOfBeingErnest 的答案更快或更慢。

Answer 2

这里似乎不需要stem 图，因为标记无论如何都是不可见的，并且由于点数很多而没有意义。

使用 LineCollection 可能更有意义。无论如何，这是how matplotlib will do it in a future version - 请参阅this PR。下面的代码对我来说在 0.25 秒内运行。 （这仍然比使用plot 稍长，因为行数较多。）

import numpy as np
from scipy.fftpack import fft
import matplotlib.pyplot as plt
import time
import matplotlib.collections as mcoll

N=2048
k = np.arange(0,N)
dr = 100

cos = np.cos
pi = np.pi

w = 1-1.932617*cos(2*pi*k/(N-1))+1.286133*cos(4*pi*k/(N-1))-0.387695*cos(6*pi*k/(N-1))+0.0322227*cos(8*pi*k/(N-1))

y = np.concatenate([w, np.zeros((7*N))])

H = abs(fft(y, axis = 0))
H = np.fft.fftshift(H)
H = H/max(H)
H = 20*np.log10(H)
H = dr+H 
H[H < 0] = 0        # Set all negative values in dr+H to 0

plotdata = ((np.arange(1,(8*N)+1,1))-1-4*N)/8


lines = []
for thisx, thisy in zip(plotdata,H):
    lines.append(((thisx, 0), (thisx, thisy)))
stemlines = mcoll.LineCollection(lines, linestyles="-",
                    colors="C0", label='_nolegend_')
plt.gca().add_collection(stemlines)


plt.axis([(-4*N)/8, (4*N)/8, 0, dr])    
plt.grid()
plt.ylabel('decibels')
plt.xlabel('DFT bins')
plt.title('Frequency response (Flat top)')

plt.show()