PSD additions

obspy/signal/spectral_estimation.py

@@ -967,6 +967,228 @@ def get_NHNM():
     return (periods, nlnm)
 
 
+def plot_psd_trace(self, plot_type='default', nperseg=4096, noverlap=None,
+                   detrend=mlab.detrend_linear, window=mlab.window_hanning,
+                   convert_to_periods=False, convert_to_db=False, logx=False,
+                   plot_noise_models=False, ax=None, show_plot=True,
+                   amplitude_label_units=None, calc_only=False):
+    """
+    Plot the power spectral density estimated using Welch's method.
+
+    Welch's method works by averaging modified periodograms calculated from
+    overlapping data segments (see also [McNamara2004]_). When used without
+    windowing and with zero overlap, this is identical to Bartlett's method.
+
+    Setting plot_type to 'ppsd' sets up the plot very similar to the plot
+    generated by obspy.signal.spectral_estimation.PPSD. Note that in this mode
+    there the visible area contains data if the waveform data was converted to
+    ground acceleration.
+
+    :type plot_type: String
+    :param plot_type: XXX move to special method?
+    :type nperseg: int
+    :param nperseg: Number of data points to use for FFT. Care must be taken
+            to choose a number suitable for FFT. See e.g.
+            :func:`~obspy.signal.util.prevpow2`.
+    :type noverlap: int
+    :param noverlap: The number of overlapping points. If `None`, defaults to
+            50% of nperseg.
+    XXX Document other options (if they stay).
+    """
+    if plot_type == 'ppsd':
+        # This is the same as in obspy.signal.spectral_estimation.PPSD.
+        # Initially split the trace into 13 segments, overlapping by 75%
+        segment_length = self.stats.npts / 4.0
+        # Reduce segment length to next lower power of 2 (potentially
+        # increasing number of segments up to a maximum of 27)
+        nperseg = prevpow2(segment_length)
+        noverlap = int(0.75 * nperseg)
+        # Since plot_type is merely a shortcut, force appropriate options,
+        # ignoring manually set values
+        detrend = mlab.detrend_linear
+        window = fft_taper
+        convert_to_periods = True
+        convert_to_db = True
+        logx = True
+        plot_noise_models = True
+
+    if noverlap is None:
+        # Default to 50% as appropriate for window_hanning
+        noverlap = int(0.5 * nperseg)
+
+    # Perform the spectral estimation
+    Pxx, f = psd(
+        self.data,
+        NFFT=nperseg,
+        Fs=self.stats.sampling_rate,
+        detrend=detrend,
+        window=window,
+        noverlap=noverlap
+    )
+
+    # Remove first term (zero frequency)
+    Pxx = Pxx[1:]
+    f = f[1:]
+
+    if calc_only:
+        return Pxx, f
+
+    if convert_to_periods:
+        # Go from frequency to period
+        f = 1. / f
+
+    if convert_to_db:
+        # Replace zero values for safe logarithm (as in PPSD)
+        dtiny = np.finfo(0.0).tiny
+        idx = Pxx < dtiny
+        Pxx[idx] = dtiny
+        # Go to db
+        Pxx = np.log10(Pxx)
+        Pxx *= 10
+
+    if ax is None:
+        _fig, ax = plt.subplots(nrows=1)
+
+    ax.plot(f, Pxx)
+
+    if plot_noise_models:
+        # Taken from PPSD.plot()
+        model_periods, high_noise = get_NHNM()
+        ax.plot(model_periods, high_noise, '0.4', linewidth=2)
+        model_periods, low_noise = get_NLNM()
+        ax.plot(model_periods, low_noise, '0.4', linewidth=2)
+
+    if plot_type == 'ppsd':
+        # Set the same limits as in PPSD with 1h segments
+        ax.set_xlim(0.01, 179)
+        ax.set_ylim(-200, -50)
+        # Warn if there is nothing to display (likely data wasn't acceleration)
+        if not np.any((Pxx > -200) & (Pxx < -50)):
+            msg = 'No data to display (PSD data is between %.01f and ' + \
+                  '%.01f dB). Maybe instrument response was not removed ' + \
+                  'or data is not ground acceleration?'
+            warnings.warn(msg % (Pxx.min(), Pxx.max()))
+    else:
+        # Scale to maximum of central 90% (so that scaling won't be dominated
+        # by border effects)
+        cut = int(0.05 * len(Pxx))
+        y_max = np.max(Pxx[cut:-cut])
+        # Add some padding at the top
+        y_max *= 1.05
+        ax.set_ylim(0, y_max)
+
+    if logx:
+        ax.semilogx()
+
+    # Set plot labels
+    if convert_to_periods:
+        ax.set_xlabel('Period [s]')
+    else:
+        ax.set_xlabel('Frequency [Hz]')
+    if convert_to_db:
+        ax.set_ylabel('Amplitude [dB]')
+    else:
+        label = 'Amplitude'
+        if amplitude_label_units:
+            label += ' [' + amplitude_label_units + ']'
+        ax.set_ylabel(label)
+    ax.xaxis.set_major_formatter(FormatStrFormatter("%.2f"))
+    title = "%s   PSD Estimation"
+    ax.set_title(title % self.id)
+
+    if show_plot:
+        plt.show()
+
+
+def plot_psd_stream(self, **kwargs):
+    """
+    Same as plot_psd_trace, but for Stream objects.
+    """
+    self.merge(fill_value=0)
+
+    n_traces = len(self)
+
+    if not n_traces:
+        msg = 'No traces in stream'
+        raise IndexError(msg)
+    elif 'ax' in kwargs:
+        # Catch keyword ax, which is explicitly provided later, and silently
+        # convert to axes (if axes was not specified)
+        # XXX probably better to always raise
+        if n_traces != 1 or 'axes' in kwargs:
+            msg = "Unexpected keyword argument 'ax'"
+            raise TypeError(msg)
+        axes = (kwargs['ax'],)
+    elif 'axes' in kwargs:
+        # Check if right number of axes was provided in keywords arguments
+        if len(kwargs['axes']) != n_traces:
+            msg = 'Number of axes must match number of traces after merging'
+            raise ValueError(msg)
+        axes = kwargs['axes']
+    elif n_traces == 1:
+        _fig, ax0 = plt.subplots(nrows=1)
+        axes = (ax0,)
+    else:
+        _fig, axes = plt.subplots(nrows=n_traces)
+        plt.subplots_adjust(hspace=1.2)
+
+    largest_ymax = 0
+    for trace, ax in zip(self, axes):
+        kwargs['ax'] = ax
+        kwargs['show_plot'] = False
+        trace.plot_psd(**kwargs)
+        ymax = ax.get_ylim()[1]
+        if ymax > largest_ymax:
+            largest_ymax = ymax
+
+    if not kwargs.get('plot_type') == 'ppsd':
+        for ax in axes:
+            ax.set_ylim(0, largest_ymax)
+
+    plt.show()
+
+
+def plot_psd_mtspec_trace(self, ax=None, show_plot=True, **kwargs):
+    """
+    Minimal method to plot the spectrum generated using mtspec for quick
+    comparison.
+    """
+    import mtspec
+    Pxx, f = mtspec.mtspec(self.data, self.stats.delta, 4)
+
+    if ax is None:
+        _fig, ax = plt.subplots(nrows=1)
+
+    ax.plot(f, Pxx)
+
+    # Scale to maximum of central 90% (so that scaling won't be dominated
+    # by border effects)
+    cut = int(0.05 * len(Pxx))
+    y_max = np.max(Pxx[cut:-cut])
+    # Add some padding at the top
+    y_max *= 1.05
+    ax.set_ylim(0, y_max)
+
+    ax.set_xlabel('Frequency [Hz]')
+    ax.set_ylabel('Amplitude [??]')
+    ax.xaxis.set_major_formatter(FormatStrFormatter("%.2f"))
+    title = "%s   PSD Estimation (mtspec)"
+    ax.set_title(title % self.id)
+
+    if show_plot:
+        plt.show()
+
+
+def __monkey_patch():
+    Stream.plot_psd = plot_psd_stream
+    Trace.plot_psd = plot_psd_trace
+    Trace.plot_psd_mtspec = plot_psd_mtspec_trace
+
+
+# Apply the monkey path when the module is imported.
+__monkey_patch()
+
+
 if __name__ == '__main__':
     import doctest
     doctest.testmod(exclude_empty=True)