Spectrogram for seismology based on obspy.imaging.spectrogram

spectrogram.py

### spectrogram.py
### Plot detailed spectrogram for seismology
###     based on `obspy.imaging.spectrogram`.
###
### Last editted: Aug 09, 2020


import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from matplotlib.ticker import ScalarFormatter

import obspy as ob
import obspy.imaging.spectrogram
from obspy.imaging.cm import obspy_sequential

class PlotSpectrogram():
    """
    Plot detailed spectrogram.

    Parameters
    ----------
    tr: obspy.Trace 
        trace to plot spectrogram.
    figsize: (float, float), optional, default: None
        Total figsize as (width, height) in inches.

    Returns
    -------
    plotspectrogram: `spectrogram.PlotSpectrogram`
        The `.PlotSpectrogram` instance will be returned. 
        The original trace is stored in the attribute `tr`.

    Usage
    -----
    >>> import obspy as ob
    >>> from spectrogram import PlotSpectrogram
    >>> test_tr = ob.read('http://examples.obspy.org/RJOB_061005_072159.ehz.new')[0]
    >>> ps = PlotSpectrogram(tr=test_tr, figsize=[12,9])
    >>> ps.plot(per_lap=0.99)
    >>> ps.set(yscale='log', xlim=(80,120), ylim=(0.1,40))

    You can modify the base figure setting using `plt.rcParams`.

    Notes
    -----
    You can set the additional `Figure` parameters as kwargs.

    """
    def __init__(self, tr:ob.Trace, figsize=None, **kwargs):
        if type(tr) != ob.Trace:
            raise AttributeError(
                "`tr` should be an obspy.Trace object."
            )
        self.tr = tr
        self._fig_init(figsize=figsize, **kwargs)

    def __str__(self):
        return self.tr.__str__()
    
    def _fig_init(self, figsize=None, **kwargs):
        """
        Init figure. This is called automatically.
        """
        ## Init Fig and GridSpec
        self.fig = plt.figure(figsize=figsize, **kwargs)
        self.gs_master = GridSpec(nrows=4,ncols=20)
        
        ## Axes for waveform
        self.axes_wav = self.fig.add_subplot(
            self.gs_master[0,:-1], ylabel='Velocity'+'\n'+'[nm/s]',
        )
        self.axes_wav.tick_params(
            axis='x', direction='in', 
            labelbottom=True, bottom=True, top=True,
        )
        self.axes_wav.yaxis.set_major_formatter(ScalarFormatter(useMathText=True))
        self.axes_wav.ticklabel_format(style="sci",  axis="y",scilimits=(-1,0))
        
        ## Axes for spectrogram
        self.axes_main = self.fig.add_subplot(
            self.gs_master[1:,:-1], sharex=self.axes_wav,
            xlabel='Time [s]',ylabel='Frequency [Hz]',
        )
        self.axes_main.tick_params(
            axis='both', direction='out', labeltop=False, top=True)
        
        ## Axes for colorbar
        self.axes_color = self.fig.add_subplot(
            self.gs_master[1:,-1],xlabel='PSD',
        )
        self.axes_color.tick_params(
            axis='y', direction='out', labelright=True, 
            labelleft=False,left=False,right=True)
        
        self.fig.suptitle(
            t=f"{self.tr.stats.station}.{self.tr.stats.channel} {self.tr.stats.starttime}", y=0.93,
        )
        self.fig.subplots_adjust(hspace=0.3)
        self.fig.subplots_adjust(wspace=1.5)
    
    def plot(self, per_lap=0.9, wlen=None, cmap=obspy_sequential, **kwargs):
        """
        Plot spectrogram using obspy.imaging.spectrogram.spectrogram()

        Parameters
        ----------
        per_lap : float, default to 0.9
            Percentage of overlap of sliding window, ranging from 0 to 1. High overlaps take a long time to compute.

        wlen : int or float
            Window length for fft in seconds. If this parameter is too small, the calculation will take forever. If None, it defaults to (samp_rate/100.0).

        cmap : `~matplotlib.colors.Colormap`
            Specify a custom colormap instance. If not specified, then the default ObsPy sequential colormap is used.
        
        Notes
        -----
        A bug in obspy.imaging.spectrogram.spectrogram for `log=True`, which does not return mappable images.
        For the `log` scale, use
            set(yscale='log')
        """
        self.axes_wav.plot(self.tr.times(),self.tr.data, color='k', lw=0.8)
        ob.imaging.spectrogram.spectrogram(
            data=self.tr.data, samp_rate=self.tr.stats.sampling_rate, 
            per_lap=per_lap, wlen=wlen, cmap=cmap, log=False, 
            axes=self.axes_main, show=False, **kwargs)
        mappable = self.axes_main.images[0]
        plt.colorbar(mappable=mappable, cax=self.axes_color)
        
    def set(self, yscale='linear', **kwargs):
        """
        Figure settings.
        
        Parameters
        ----------
        yscale: 'log' or 'linear', default to 'linear'

        For other params, see https://matplotlib.org/3.2.2/api/_as_gen/matplotlib.axes.Axes.set.html
        """
        self.axes_main.set(yscale=yscale, **kwargs)



################ Module status ######################
__author__ = "Yasu Sawaki"
__status__ = "production"
__date__ = "Aug 09, 2020"
__version__ = "0.1.1"
__doc__ = PlotSpectrogram.__doc__
__all__ = [
    'PlotSpectrogram',
    # '__author__', '__status__', '__date__', '__version__'
]
################ Module status ######################