Source code for bruges.attribute.spectrogram
# -*- coding: utf-8 -*-
"""
Spectrogram.
:copyright: 2015 Agile Geoscience
:license: Apache 2.0
"""
import numpy as np
from scipy.fftpack import fft
from scipy.signal import get_window
from bruges.util import next_pow2
[docs]def spectrogram(data, window_length,
dt=1.0,
window_type='boxcar',
overlap=0.5,
normalize=False,
zero_padding=0):
"""
Calculates a spectrogram using windowed STFTs.
Args:
data (array): 1D numpy array to process into spectra.
window_length (float): The length of the window in seconds if dt is set, otherwise in samples. Will zero pad to the closest power of two.
window_type: A string specifying the type of window to use for the STFT. The same input as scipy.signal.get_window
dt (float or int): The time sample interval of the trace. Defaults to 1, which allows window_length to be in seconds.
overlap (float): The fractional overlap for each window. A value of 0 uses no redudant data, a value of 1 slides the STFT window one sample at a time. Defaults to 0.5
normalize (bool): Normalizes the each spectral slice to have unity energy.
zero_padding (float or int): The amount of zero padding to when creating the spectra.
Returns
-------
out : A spectrogram of the data ([time, freq]). (2D array for 1D input)
See Also
--------
spectraldecomp : Spectral decomposition
"""
# Make the base window
window_n = int(np.floor(window_length / dt))
pad = int(np.floor(zero_padding / dt))
window = get_window(window_type, window_n)
# Calculate how many STFTs we need to do.
stride = int(window.size * (1 - overlap) + 1)
n_windows = int(np.ceil((data.size - window.size) / stride) + 1)
# Pad window to power of 2
padded_window = np.zeros(next_pow2(window.size+pad))
# Init the output array
output = np.zeros([n_windows, int(padded_window.size // 2)])
# Do the loop
for i in range(int(n_windows)):
start = int(i * stride)
end = start + int(window.size)
# Do not compute last samples if we don't have a full window
if (end > data.size-1):
break
padded_window[0:window.size] = window*data[start:end]
spect = (2. * np.absolute(fft(padded_window)) /
window.size)[0:int(padded_window.size // 2)]
if normalize:
output[i, :] = spect / np.sum(spect**2)
else:
output[i, :] = spect
return output