1# -*- coding: utf-8 -*- 2# ------------------------------------------------------------------ 3# Filename: spectrogram.py 4# Purpose: Plotting spectrogram of Seismograms. 5# Author: Christian Sippl, Moritz Beyreuther 6# Email: sippl@geophysik.uni-muenchen.de 7# 8# Copyright (C) 2008-2012 Christian Sippl 9# -------------------------------------------------------------------- 10""" 11Plotting spectrogram of seismograms. 12 13:copyright: 14 The ObsPy Development Team (devs@obspy.org) 15:license: 16 GNU Lesser General Public License, Version 3 17 (https://www.gnu.org/copyleft/lesser.html) 18""" 19from __future__ import (absolute_import, division, print_function, 20 unicode_literals) 21from future.builtins import * # NOQA @UnusedWildImport 22 23import math 24 25import numpy as np 26from matplotlib import mlab 27from matplotlib.colors import Normalize 28 29from obspy.imaging.cm import obspy_sequential 30 31 32def _nearest_pow_2(x): 33 """ 34 Find power of two nearest to x 35 36 >>> _nearest_pow_2(3) 37 2.0 38 >>> _nearest_pow_2(15) 39 16.0 40 41 :type x: float 42 :param x: Number 43 :rtype: Int 44 :return: Nearest power of 2 to x 45 """ 46 a = math.pow(2, math.ceil(np.log2(x))) 47 b = math.pow(2, math.floor(np.log2(x))) 48 if abs(a - x) < abs(b - x): 49 return a 50 else: 51 return b 52 53 54def spectrogram(data, samp_rate, per_lap=0.9, wlen=None, log=False, 55 outfile=None, fmt=None, axes=None, dbscale=False, 56 mult=8.0, cmap=obspy_sequential, zorder=None, title=None, 57 show=True, sphinx=False, clip=[0.0, 1.0]): 58 """ 59 Computes and plots spectrogram of the input data. 60 61 :param data: Input data 62 :type samp_rate: float 63 :param samp_rate: Samplerate in Hz 64 :type per_lap: float 65 :param per_lap: Percentage of overlap of sliding window, ranging from 0 66 to 1. High overlaps take a long time to compute. 67 :type wlen: int or float 68 :param wlen: Window length for fft in seconds. If this parameter is too 69 small, the calculation will take forever. If None, it defaults to 70 (samp_rate/100.0). 71 :type log: bool 72 :param log: Logarithmic frequency axis if True, linear frequency axis 73 otherwise. 74 :type outfile: str 75 :param outfile: String for the filename of output file, if None 76 interactive plotting is activated. 77 :type fmt: str 78 :param fmt: Format of image to save 79 :type axes: :class:`matplotlib.axes.Axes` 80 :param axes: Plot into given axes, this deactivates the fmt and 81 outfile option. 82 :type dbscale: bool 83 :param dbscale: If True 10 * log10 of color values is taken, if False the 84 sqrt is taken. 85 :type mult: float 86 :param mult: Pad zeros to length mult * wlen. This will make the 87 spectrogram smoother. 88 :type cmap: :class:`matplotlib.colors.Colormap` 89 :param cmap: Specify a custom colormap instance. If not specified, then the 90 default ObsPy sequential colormap is used. 91 :type zorder: float 92 :param zorder: Specify the zorder of the plot. Only of importance if other 93 plots in the same axes are executed. 94 :type title: str 95 :param title: Set the plot title 96 :type show: bool 97 :param show: Do not call `plt.show()` at end of routine. That way, further 98 modifications can be done to the figure before showing it. 99 :type sphinx: bool 100 :param sphinx: Internal flag used for API doc generation, default False 101 :type clip: [float, float] 102 :param clip: adjust colormap to clip at lower and/or upper end. The given 103 percentages of the amplitude range (linear or logarithmic depending 104 on option `dbscale`) are clipped. 105 """ 106 import matplotlib.pyplot as plt 107 # enforce float for samp_rate 108 samp_rate = float(samp_rate) 109 110 # set wlen from samp_rate if not specified otherwise 111 if not wlen: 112 wlen = samp_rate / 100. 113 114 npts = len(data) 115 # nfft needs to be an integer, otherwise a deprecation will be raised 116 # XXX add condition for too many windows => calculation takes for ever 117 nfft = int(_nearest_pow_2(wlen * samp_rate)) 118 if nfft > npts: 119 nfft = int(_nearest_pow_2(npts / 8.0)) 120 121 if mult is not None: 122 mult = int(_nearest_pow_2(mult)) 123 mult = mult * nfft 124 nlap = int(nfft * float(per_lap)) 125 126 data = data - data.mean() 127 end = npts / samp_rate 128 129 # Here we call not plt.specgram as this already produces a plot 130 # matplotlib.mlab.specgram should be faster as it computes only the 131 # arrays 132 # XXX mlab.specgram uses fft, would be better and faster use rfft 133 specgram, freq, time = mlab.specgram(data, Fs=samp_rate, NFFT=nfft, 134 pad_to=mult, noverlap=nlap) 135 # db scale and remove zero/offset for amplitude 136 if dbscale: 137 specgram = 10 * np.log10(specgram[1:, :]) 138 else: 139 specgram = np.sqrt(specgram[1:, :]) 140 freq = freq[1:] 141 142 vmin, vmax = clip 143 if vmin < 0 or vmax > 1 or vmin >= vmax: 144 msg = "Invalid parameters for clip option." 145 raise ValueError(msg) 146 _range = float(specgram.max() - specgram.min()) 147 vmin = specgram.min() + vmin * _range 148 vmax = specgram.min() + vmax * _range 149 norm = Normalize(vmin, vmax, clip=True) 150 151 if not axes: 152 fig = plt.figure() 153 ax = fig.add_subplot(111) 154 else: 155 ax = axes 156 157 # calculate half bin width 158 halfbin_time = (time[1] - time[0]) / 2.0 159 halfbin_freq = (freq[1] - freq[0]) / 2.0 160 161 # argument None is not allowed for kwargs on matplotlib python 3.3 162 kwargs = {k: v for k, v in (('cmap', cmap), ('zorder', zorder)) 163 if v is not None} 164 165 if log: 166 # pcolor expects one bin more at the right end 167 freq = np.concatenate((freq, [freq[-1] + 2 * halfbin_freq])) 168 time = np.concatenate((time, [time[-1] + 2 * halfbin_time])) 169 # center bin 170 time -= halfbin_time 171 freq -= halfbin_freq 172 # Log scaling for frequency values (y-axis) 173 ax.set_yscale('log') 174 # Plot times 175 ax.pcolormesh(time, freq, specgram, norm=norm, **kwargs) 176 else: 177 # this method is much much faster! 178 specgram = np.flipud(specgram) 179 # center bin 180 extent = (time[0] - halfbin_time, time[-1] + halfbin_time, 181 freq[0] - halfbin_freq, freq[-1] + halfbin_freq) 182 ax.imshow(specgram, interpolation="nearest", extent=extent, **kwargs) 183 184 # set correct way of axis, whitespace before and after with window 185 # length 186 ax.axis('tight') 187 ax.set_xlim(0, end) 188 ax.grid(False) 189 190 if axes: 191 return ax 192 193 ax.set_xlabel('Time [s]') 194 ax.set_ylabel('Frequency [Hz]') 195 if title: 196 ax.set_title(title) 197 198 if not sphinx: 199 # ignoring all NumPy warnings during plot 200 with np.errstate(all='ignore'): 201 plt.draw() 202 if outfile: 203 if fmt: 204 fig.savefig(outfile, format=fmt) 205 else: 206 fig.savefig(outfile) 207 elif show: 208 plt.show() 209 else: 210 return fig 211