from IPython.core.display import Image 
Image(url='http://labrosa.ee.columbia.edu/crucialpython/logo.png', width=600) 

# Crucial imports
import numpy as np
import matplotlib.pyplot as plt

# Let's say we have a (near) periodic signal
x = np.sin(np.arange(128))
plt.plot(x)

# Build a "framed" version of x as successive, overlapped sequences
# of frame_length points
from numpy.lib import stride_tricks
frame_length = 16
hop_length = 4
num_frames = 1 + (len(x) - frame_length) / hop_length
row_stride = x.itemsize * hop_length
col_stride = x.itemsize
x_framed = stride_tricks.as_strided(x, 
                                    shape=(num_frames, frame_length), 
                                    strides=(row_stride, col_stride))
plt.imshow(x_framed, interpolation='nearest', cmap='gray')

# If we take the FFT of each row, we can see the short-time fourier transform
plt.imshow(np.abs(np.fft.rfft(x_framed)), interpolation='nearest')

# Although there's a steady sinusoidal component, we see interference between the 
# window frame and the signal phase.  We need a tapered window applied to each frame.
window = np.hanning(frame_length)
plt.plot(window)

# But what's the best way to multiply each frame of x_framed by window?
# Linear algebra way is to multiply by a diagonal matrix
diag_window = np.diag(window)
plt.imshow(diag_window, interpolation='nearest', cmap='gray')

# Now apply it to each frame using matrix multiplication
x_framed_windowed = np.dot(x_framed, diag_window)
plt.imshow(x_framed_windowed, interpolation='nearest', cmap='gray')

# Matlab way is to construct a matrix of repeating rows of the same size
window_repeated = np.tile(window, (num_frames,1))
plt.imshow(window_repeated, interpolation='nearest', cmap='gray')
# then pointwise multiplication applies it to each row
x_framed_windowed = x_framed * window_repeated

# Numpy broadcasting implicitly (and efficiently) repeats singleton or missing dimensions
# to make matrices the same size, so tiling is unneeded
plt.imshow(x_framed*window, interpolation='nearest', cmap='gray')

# Compare the timings:
%timeit np.dot(x_framed, np.diag(window))
%timeit x_framed*np.tile(window, (num_frames,1))
%timeit x_framed*window
# The big win is not having to allocate the second num_frames x frame_length array

# What about if we had our data in columns?
x_framed_T = x_framed.T
plt.imshow(x_framed_T * window, interpolation='nearest', cmap='gray')

# Broadcast works by starting at the last dimensions (the fastest-changing ones in 'C' ordering) 
# and promoting either one if it's one.  
# So we just have to make our window be frame_length x 1
# We can do this with slicing and np.newaxis:
plt.imshow(x_framed_T * window[:,np.newaxis], interpolation='nearest', cmap='gray')
# Now broadcasting works again

# Broadcasting works across multiple dimensions.  
# It goes through each dimension from the last, looking for a match 
# or promoting singletons
a = np.random.rand( 3, 4, 5 )
b = np.random.rand( 3, 5 )
c = a*b

# That didn't work because there was no singleton dimension to promote
# so we can introduce one, with reshape for instance:
b2 = np.reshape(b, (3, 1, 5))
c1 = a*b2
# or using slicing
c2 = a * b[:, np.newaxis, :]
print np.allclose(c1, c2)

# Remember why we wanted the windowing, to remove artefacts from the STFT?
plt.subplot(121)
plt.imshow(np.abs(np.fft.rfft(x_framed)), interpolation='nearest')
plt.subplot(122)
plt.imshow(np.abs(np.fft.rfft(x_framed * window)), interpolation='nearest')
# Windowing drastically reduces framing-related artefacts 
# (at the cost of a little frequency resolution)