Estou tentando produzir um gráfico Stem usando a função 'matplotlib.pyplot.stem'. O código funciona, mas leva mais de 5 minutos para ser processado.

Eu tenho um código semelhante no Matlab que produz o mesmo gráfico com os mesmos dados de entrada quase instantaneamente.

Existe uma maneira de otimizar esse código para velocidade ou uma função melhor que eu poderia estar usando?

Os argumentos para o gráfico de tronco 'H' e 'plotdata' são matrizes 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

Aqui também está o código do Matlab para referência:

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