m-thesis-introduction/lib/ft_plot.py

"""
Functions to simplify plotting of fourier spectra
"""
# vim: fdm=indent ts=4

import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np

def plot_amplitude(*args, **kwargs):
    kwargs = { 'plot_amplitude': True, **kwargs}
    return plot_spectrum(*args, **kwargs)

def plot_spectrum(
        spectrum, freqs,
        plot_complex=False, plot_power=False, plot_amplitude=None,
        freq_unit="Hz", freq_scaler=1,
        title='Spectrum', xlabel='Frequency', ax=None,
        ylabel=True, grid=True,
        **plot_kwargs
        ):
    """ Plot a signal's spectrum on an Axis object"""
    plot_amplitude = plot_amplitude or (not plot_power and not plot_complex)
    alpha = 1

    if ax is None:
        ax = plt.gca()

    if title:
        ax.set_title(title)
    if xlabel:
        ax.set_xlabel(xlabel + (" ["+freq_unit+"]" if freq_unit else "" ))
    if ylabel is True:
        ylabel = ""
        if plot_amplitude or plot_complex:
            ylabel = "Amplitude"
        if plot_power:
            if ylabel:
                ylabel += "|"
            ylabel += "Power"
    if ylabel:
        ax.set_ylabel(ylabel)

    ax.grid(grid)
    
    plot_kwargs = {'marker': '.', 'linestyle': 'solid', **plot_kwargs}
    if plot_complex:
        alpha = 0.5
        ax.plot(freqs/freq_scaler, np.real(spectrum),   **{**dict(label='Real', alpha=alpha, ), **plot_kwargs})
        ax.plot(freqs/freq_scaler, np.imag(spectrum),   **{**dict(label='Imag', alpha=alpha, ), **plot_kwargs})

    if plot_power:
        ax.plot(freqs/freq_scaler, np.abs(spectrum)**2, **{**dict(label='Power', alpha=alpha,), **plot_kwargs})

    if plot_amplitude:
        ax.plot(freqs/freq_scaler, np.abs(spectrum),    **{**dict(label='Abs', alpha=alpha,  ), **plot_kwargs})

    ax.legend()

    return ax

def plot_phase(
        spectrum, freqs,
        ylim_epsilon=0.5, ax=None, grid=True,
        freq_unit="Hz", freq_scaler=1, xlabel='Frequency',
        ylabel="Phase",
        major_divider=2, minor_divider=12,
        **plot_kwargs
    ):
    """
    Plot the phase of spectrum
    """

    if ax is None:
        ax = plt.gca()

    if ylabel:
        ax.set_ylabel("Phase")
    if xlabel:
        ax.set_xlabel(xlabel + (" ["+freq_unit+"]" if freq_unit else "" ))

    ax.grid(grid)
    
    plot_kwargs = {'marker': '.', 'linestyle': 'solid', **plot_kwargs}
    ax.plot(freqs/freq_scaler, np.angle(spectrum), **plot_kwargs)
    ax.set_ylim(-1*np.pi - ylim_epsilon, np.pi + ylim_epsilon)

    axis_pi_ticker(ax.yaxis, major_divider=major_divider, minor_divider=minor_divider)

    return ax

def axis_pi_ticker(axis, major_divider=2, minor_divider=12):

    major_tick = Multiple(major_divider, np.pi, '\pi')
    minor_tick = Multiple(minor_divider, np.pi, '\pi')

    axis.set_major_locator(major_tick.locator())
    axis.set_major_formatter(major_tick.formatter())
    axis.set_minor_locator(minor_tick.locator())

    return axis

def plot_signal(
        signal, sample_rate = 1,
        time=None, ax=None,
        title='Signal', ylabel='Amplitude',
        time_unit="s", xlabel='Time',
        **kwargs):
    """
    Plot the signal in the time domain
    """

    if ax is None:
        ax = plt.gca()

    if time is None:
        time = np.arange(len(signal))/sample_rate

    if title:
        ax.set_title(title)
    if xlabel:
        ax.set_xlabel(xlabel + (" ["+time_unit+"]" if time_unit else "" ))
    if ylabel:
        ax.set_ylabel(ylabel)

    ax.plot(time, signal, **kwargs)

    return ax

def combined_spectrum_figure(
        spectrum, freqs,
        fig=None, gs=None,
        spectrum_ax=None, phase_ax=None,
        spectrum_kwargs={}, phase_kwargs={},
        **shared_kwargs
        ):
    """
    Plot both the frequencies and phase in one figure.
    """

    ax1, ax2 = None, None
    if spectrum_ax is not None or phase_ax is not None:
        if spectrum_ax is not None:
            ax1 = spectrum_ax
        if phase_ax is not None:
            ax2 = phase_ax

    else:
        # configure plotting layout
        if fig is None:
            fig = plt.figure()

        if gs is None:
            gs = gridspec.GridSpec(2, 1, figure=fig, height_ratios=[3,1], hspace=0)

        ax1 = fig.add_subplot(gs[:-1, -1])
        ax2 = fig.add_subplot(gs[-1, -1], sharex=ax1)

    axes = np.array([ax1, ax2])

    # plot the spectrum
    plot_spectrum(spectrum, freqs, ax=ax1, **{ **shared_kwargs, **spectrum_kwargs})

    # plot the phase
    plot_phase(spectrum, freqs, ax=ax2, **{ **shared_kwargs, **phase_kwargs})

    ax1.xaxis.tick_top()
    [label.set_visible(False) for label in ax1.get_xticklabels()]

    return fig

def multiple_formatter(denominator=2, number=np.pi, latex='\pi'):
    """
    From https://stackoverflow.com/a/53586826
    """
    def gcd(a, b):
        while b:
            a, b = b, a%b
        return a

    def _multiple_formatter(x, pos):
        den = denominator
        num = np.int(np.rint(den*x/number))
        com = gcd(num,den)
        (num,den) = (int(num/com),int(den/com))
        if den==1:
            if num==0:
                return r'$0$'
            if num==1:
                return r'$%s$'%latex
            elif num==-1:
                return r'$-%s$'%latex
            else:
                return r'$%s%s$'%(num,latex)
        else:
            if num==1:
                return r'$\frac{%s}{%s}$'%(latex,den)
            elif num==-1:
                return r'$\frac{-%s}{%s}$'%(latex,den)
            else:
                return r'$\frac{%s%s}{%s}$'%(num,latex,den)
    return _multiple_formatter

class Multiple:
    """
    From https://stackoverflow.com/a/53586826
    """
    def __init__(self, denominator=2, number=np.pi, latex='\pi'):
        self.denominator = denominator
        self.number = number
        self.latex = latex

    def locator(self):
        return plt.MultipleLocator(self.number / self.denominator)

    def formatter(self):
        return plt.FuncFormatter(multiple_formatter(self.denominator, self.number, self.latex))