Source

stockwelltransform / stockwell / plots.py

Full commit
import matplotlib.pyplot as plt

from matplotlib.offsetbox import TextArea, DrawingArea, OffsetImage, \
     AnnotationBbox

from mpl_toolkits.axes_grid1 import host_subplot

def plotspec(psx, fs=2.0, lofreq=None, hifreq=None, t1=None, t2=None):
    """
    useful for plotting the power of a stockwell transform
    it relies upon matplotlib for display
    example:
    # for a signal x, with sampling frequency 200
    >>> import stockwell, pylab
    >>> import stockwell.plots as plots
    >>> x = pylab.zeros(1000.0) # 5 seconds
    >>> fs = 200 # sample rate
    >>> x[250:350] = 1.0 # step function
    >>> sx = stockwell.st(x)
    >>> psx = abs(sx) # create power
    >>> r=plots.plotspec(psx,200)
    >>> # pylab.show() # to visualize this
    """
    extent = [0,psx.shape[1], 0.0, fs/2.0] # default extents
    if t1 != None and t2 != None:
        extent[0] = t1
        extent[1] = t2
    if lofreq != None:
        extent[2] = lofreq
    if hifreq != None:
        extent[3] = hifreq
    
    return plt.imshow(psx, extent=extent, aspect='auto', origin='lower')



def stpowerstack(x,stx):
    """
    need to add row labels
    """
    ax1=host_subplot(211)
    plt.plot(x)
    ax2 = host_subplot(212)
    #pax2 = ax2.twinx()
    #pax2.set_ylabel('frequency(Hz)')
    plt.imshow(abs(stx), aspect='auto')
    plt.ylabel('frequency(Hz)')
    ylocs,ylabels = yticks()
    #plt.ylabel('st-row(f*L/fs)')
    yt,xt = stx.shape


    return ax2
    
if __name__ == "__main__":
    import doctest
    doctest.testmod()