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