matplotlib茎图的优化 [英] Optimization of matplotlib stem plot

问题描述

我正在尝试使用'matplotlib.pyplot.stem'函数生成一个Stem图.该代码可以运行，但是要花5分钟以上的时间.

I'm trying to produce a Stem plot using the 'matplotlib.pyplot.stem' function. The code works but it is taking over 5 minutes to process.

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

I have a similar code within Matlab that produces the same plot with the same input data almost instantly.

有没有一种方法可以优化此代码以提高速度或更好的功能?

Is there a way to optimize this code for speed or a better function I could be using?

茎图'H'和'plotdata'的参数为16384 x 1数组.

The arguments for the stem plot 'H' and 'plotdata' are 16384 x 1 arrays.

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代码供参考:

Here is also the Matlab code for reference:

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)')

推荐答案

这里似乎不需要stem绘图，因为无论如何标记都是看不见的，并且由于存在大量点而没有意义

There seems to be no need for a stem plot here, since the markers are anyway made invsibible and would not make sense due to the large number of points.

使用LineCollection可能更有意义.这是 matplotlib在将来的版本中将如何实现-参见- 此PR .下面的代码对我来说在0.25秒内运行. (由于行数很多，这仍然比使用plot稍长.)

Instead the use of a LineCollection may make sense. This is how matplotlib will do it in a future version anyways - see this PR. The code below runs within 0.25 seconds for me. (This is still slightly longer than using plot, due to the large number of lines.)

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()

这篇关于matplotlib茎图的优化的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！