senzotools.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@author: Martin Šlapák
"""

from datetime import timedelta
from datetime import datetime as dt
from scipy import signal
import numpy as np
import matplotlib.pyplot as plt
from scipy.fft import fftshift


def read_binary(fn, raw=False):
    """Load arrays from ``.npy`` or ``.npz`` files.

    Args:
        fn (string): Path to the file.

    Returns:
        numpy.array | dict: Loaded array | dictionary of numpy arrays
    """
    r = np.array([])
    with np.load(fn) as data:
        if raw:
            r = {}
            for k in data.files:
                r[k] = data[k]
        else:
            key = data.files[0]
            r = data[key]
    return r


def write_binary(fn, *args, **kwds):
    """Load arrays from ``.npy`` or ``.npz`` files.

    Args:
        fn (string): Path to the file.
        arr (numpy.array): Array to save.
    """
    np.savez_compressed(fn, *args, **kwds)


def read_hdf5(fn, key='arr'):
    import h5py
    with h5py.File(fn, 'r') as hf:
        return np.array(hf.get(key))

# TODO: probably candidate to be removed and replaced by using polarilog.utils.write_hdf5
def write_hdf5(fn, data, creator='senzotools', description=''):
    from polarilog.utils import write_hdf5 as _write_hdf5
    _write_hdf5(fn, data, creator=creator, description=description)


def get_time(stime='2018-09-20_16-08-44.272694'):
    """Parse datetime string (throw out microseconds).

    Args:
        stime (string): String to parse in following format: 2018-09-20_16-08-44.272694

    Returns:
        datetime: Parsed datetime.
    """
    if len(stime.split('.')) < 2:
        return dt.now()
    ymdhms, _ = stime.split('.')
    when = dt.strptime(ymdhms, '%Y-%m-%d_%H-%M-%S')
    return when


def normalize(a, maxval=1):
    """Normalize data to interval <0,maxval>.

    Args:
        a (numpy.array): An array to normalize.
        maxval (float): An upper bound of interval for normalization.

    Returns:
        numpy.array: An array with normalized values.
    """
    return (a - np.min(a)) * (maxval / np.max(a - np.min(a)))


def smooth(a, kernel_pts):
    """Make 1D time series data looks more smooth using convolution.

    Args:
        a (numpy.array): Array to smooth.
        kernel_pts (int): A length of convolution kernel.

    Returns:
        numpy.array: A smoothed array.
    """
    kernel = np.ones(kernel_pts) / kernel_pts
    a_smooth = np.convolve(a, kernel, mode='same')
    return a_smooth


def resample_to(data, to):
    """Resample data in array in such way that new array contains the same amount
    of values as ``to`` array.

    Args:
        data (numpy.array): An array to resample.
        to (numpy.array | int): An array to which length we need to adapt. | int lenght.

    Returns:
        numpy.array: A resampled array.
    """
    if not isinstance(to, int):
        to = to.shape[0]

    ratio = float(data.shape[0] / to)
    if ratio < 1.0:
        ratio = float(to / data.shape[0])  # up and down must be >= 1
        data = signal.resample_poly(data, up=int(ratio * 10000), down=10000)  # scales are only ints
    else:
        data = signal.resample_poly(data, up=10000, down=int(ratio * 10000))
    return data


def sum_spectrum_column(spectrum):
    """Computes the columnar sum of given spectogram.

    Args:
        spectrum (array_like): An input spectrum.

    Returns:
        numpy.array: A resulting 1D array.
    """
    return np.sum(spectrum[0:, :], axis=0)


def get_spectrum(data, fs=10, window=('tukey', 0.25), nperseg=2048, mode='psd',
                 onesided=True, noverlap=0, cut_f=None):
    # window=signal.hamming(512)
    """Plot the given spectrum or compute and plot spectrum of the given data.

    Args:
        data (array_like):
            An input spectrum.
        fs (int):
            Sampling frequency.
        window (str or tuple or array_like, optional):
            The FFT window
        nperseg (int):
            Length of each segment.
        mode (string):
            Mode of spectrogram, default = 'psd'
            available modes: ['psd', 'complex', 'magnitude', 'angle', 'phase']

    Returns:
        f : ndarray
            Array of sample frequencies.
        t : ndarray
            Array of segment times.
        Sxx : ndarray
            Spectrogram of x. By default, the last axis of Sxx corresponds
            to the segment times.
    """
    f, t, Sxx = signal.spectrogram(data, fs, window=window, nperseg=nperseg, mode=mode, return_onesided=onesided, noverlap=noverlap)
    if cut_f is not None:
        f = f[f < cut_f]
        Sxx = Sxx[0:f.shape[0], :]

    return f, t, Sxx


def plot_spectrum(data, Sxx=None, f=None, t=None, fs=10, window=('tukey', 0.25), ax=None,
                  nperseg=2048, shading='auto', upperbound=None, cmap='viridis',
                  fromtime=None, timelen=600, timestep=None, force_simple=False,
                  vmin=None, vmax=None,
                  title='Spectrogram', xlabel='Time [sec]', ylabel='Frequency [Hz]'):
    #  window=signal.hamming(512)
    """Plot the given spectrum or compute and plot spectrum of the given data.

    Args:
        data (array_like):
            An input spectrum.
        Sxx (array_like):
            Precomputed spectrogram.
        f (array_like):
            Array of freqencies (Y axis) of given spectrogram.
        t (array_like):
            Array of times (X axis) of given spectrogram.
        fs (int):
            Sampling frequency.
        window (str or tuple or array_like, optional):
            The FFT window
        ax (matplotlib.axes.Axes, optional):
            An instance of matplotlib.axes.Axes to which the spectrogram should be plotted.
        nperseg (int):
            Length of each segment.
        shading (string):
            Passed to `shading` arg of ax.pcolormesh.
        upperbound (float):
            If set, limits the Y axis from top.
        cmap (string):
            A colormap used on spectrogram-
        fromtime (datetime):
            Starting time of data.
        timelen (int):
            The time length of provided data.
        timestep (int):
            Time step for xtics labels.
        force_simple (boolean):
            When is set to False then microseconds can be displayed. Only in case
            when fromtime is set and timestep > =10.
        vmin (float, default: None):
            A lower limit for value scale, automatically computated if vmin==None.
        vmax (float, default: None):
            An upper limit for value scale, automatically computated if vmax==None.
        title (string):
            Title of spectrogram.
        xlabel (string):
            Label for X axis.
        ylabel (string):
            Label for Y axis.

    Returns
    -------
        ax : matplotlib.axes.Axes
            An instance of Axes object with plotted spectrogram.
        f : ndarray
            Array of sample frequencies.
        t : ndarray
            Array of segment times.
        Sxx : ndarray
            Spectrogram of x. By default, the last axis of Sxx corresponds
            to the segment times.
    """
    if Sxx is None:
        # modes = ['psd', 'complex', 'magnitude', 'angle', 'phase']
        f, t, Sxx = signal.spectrogram(data, fs, window=window, nperseg=nperseg, mode='psd', return_onesided=True)

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

    if timestep is None:
        timestep = timelen / 10

    if upperbound is not None:
        # ax.set_yticks(np.linspace(0, upperbound, num=9))
        # ax.set_yticklabels(f)
        f = f[f < upperbound]
        Sxx = Sxx[0:f.shape[0], :]

    # dbSxx = 10*np.log10(Sxx)
    # vmin = np.min(dbSxx)
    # vmax = np.max(dbSxx)

    # Sanitize Sxx: replace zeros with small epsilon to avoid log10(0)
    Sxx = np.where(Sxx <= 0, 1e-30, Sxx)
    im = ax.pcolormesh(t, f, 10 * np.log10(Sxx), shading=shading, cmap=cmap, vmin=vmin, vmax=vmax)
    # im = ax.pcolormesh(t, f, 10 * np.log10(Sxx), shading=shading, cmap=cmap)

    # im = ax.pcolormesh(t, f, Sxx, shading=shading, cmap=cmap)
    # ax.get_figure().colorbar(im).set_label('Intensity [dB]')
    ax.get_figure().colorbar(im, ax=ax).set_label('Intensity [dB]')
    ax.grid()

    xtickscnt = int(timelen / timestep) + 1  # +1 is for 0
    ax.set_xticks(np.linspace(t[0], t[-1], num=xtickscnt))

    if fromtime is not None:
        unit = 's'
        if timestep < 10.0:
            unit = 'ms'
            timelen = int(timelen * 1000000)
            timestep = int(timestep * 1000000)
        xlabels = []
        for i in range(0, int(timelen) + int(timestep), int(timestep)):
            if unit == 's':
                when = fromtime + timedelta(seconds=i)
                xlabels.append(str(when.strftime("%H:%M:%S")))
            else:
                when = fromtime + timedelta(microseconds=i)
                if force_simple is False:
                    xlabels.append(str(when.strftime("%H:%M:%S.%f")))
                else:
                    xlabels.append(str(when.strftime("%H:%M:%S")))

        # print('labels={}'.format(xlabels))
        # print('xtickscnt={}'.format(xtickscnt))
        # print("timestep={} tinmelen={}".format(timestep, timelen))
        ax.set_xticklabels(xlabels[0:xtickscnt], rotation=90)
    else:
        # how many decimal places should be rounded
        exp = int(f'{timelen:e}'.split('e')[-1])
        if exp > 0:
            decimal_places = 0
        else:
            decimal_places = abs(exp) + 1

        labels = [f'{i:.{decimal_places}f}' for i in np.arange(0, timelen + timestep, timestep)]
        ax.set_xticklabels(labels[0:xtickscnt], rotation=90)

    plt.ylabel(ylabel)
    plt.xlabel(xlabel)
    plt.title(title)
    return ax, f, t, Sxx


###############################################################################
# frekvence v case pres sebe
def spectrogram_sideview(cut, f, Sxx, ax, mode='db', title='Frequencies side view'):
    """
    Plots sum of spectrum over time.

    Parameters
    ----------
    cut : int
        The upper frequency limit as index into frequencies array f.
    f : np.array
        1D array of frequencies obtained from spectrogram generator.
    Sxx : np.array
        2D spectrogram array: X-time, Y-freqencies
    ax : matplotlib.axes
        Where to draw.
    mode : str, optional
        'magnitude' or 'db' defines if the values of Sxx are ploted as logarithmic (db)
        or in intensity (magnitude). The default is 'db'.

    Returns
    -------
    ax : TYPE
        DESCRIPTION.

    """
    if mode == 'magnitude':
        plt.ylabel('Intensity [mag]')
        ypos = np.linspace(np.min(Sxx[:cut, :]), np.max(Sxx[:cut, :]), 10, endpoint=True)
        for i in range(Sxx.shape[1]):
            ax.plot(f[:cut], Sxx[:cut, i], color='r', alpha=0.2)

    if mode == 'db':
        plt.ylabel('Intensity [dB]')
        ypos = np.linspace(np.min(10 * np.log10(Sxx[:cut, :])), np.max(10 * np.log10(Sxx[:cut, :])), 10, endpoint=True)
        for i in range(Sxx.shape[1]):
            ax.plot(f[:cut], 10 * np.log10(Sxx[:cut, i]), color='r', alpha=0.2)

    xpos = np.linspace(f[0], f[cut - 1], 20, endpoint=True)

    ax.set_xticks(xpos)
    ax.set_yticks(ypos)

    plt.xlabel('Frequency [Hz]')
    plt.title(title)
    return ax


def get_psd(data, fs=10, fc=10, nperseg=2048):
    """
    Computes PSD using Welch algorithm

    Parameters
    ----------
    data (array_like):
        Input signal
    fs (float):
        Sampling frequency. The default is 10.
    fc (float):
        Central frequency. The default is 10.
    nperseg: TYPE, optional
        FFT win sizw. The default is 1024.

    Returns
    -------
    f (np.array):
        Array of sample frequencies.
    Pxx (np.array 2D):
        PSD values.
    """
    onesided = False if np.iscomplexobj(data) else False

    # compute PSD by welch alg.
    f, Pxx = signal.welch(data, fs, nperseg=nperseg, return_onesided=onesided)

    # in case of complex input use fftshift to flip halves of spectrum
    if np.iscomplexobj(data):
        Pxx = fftshift(Pxx, axes=0)
        f = fftshift(f)

    # offset central frequency
    f = f + fc

    return f, Pxx


def plot_psd(data, Pxx=None, fs=10, fc=10, nperseg=1024, ax=None, semilog=False,
             title='Power Spectral Density', xlabel='frequency [Hz]',
             ylabel='PSD [V**2/Hz]'):
    """
    Plot PSD of signal.

    Parameters
    ----------
    data (array_like):
        Input signal
    Pxx (np.array):
        PSD data, if None, will be computed by get_psd(). The default is None.
    fs (float):
        Sampling frequency. The default is 10.
    fc (float):
        Central frequency. The default is 10.
    nperseg: TYPE, optional
        FFT win sizw. The default is 1024.
    ax (matplotlib.axes.Axes, optional):
        An instance of matplotlib.axes.Axes to which the spectrogram should be plotted.
    semilog (boolean):
        If should be plooted witch semi-logarithmic Y axe. The default is False.
    title (string):
        Title of spectrogram. The default is 'Power Spectral Density'.
    xlabel (string):
        Label for X axis. The default is 'frequency [Hz]'.
    ylabel (string):
        Label for Y axis. The default is 'PSD [V**2/Hz]'.

    Returns
    -------
    ax : TYPE
        DESCRIPTION.
    f : TYPE
        DESCRIPTION.
    Pxx : TYPE
        DESCRIPTION.

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

    if Pxx is None:
        f, Pxx = get_psd(data, fs=fs, fc=fc, nperseg=nperseg)

    if semilog:
        ax.semilogy(f, Pxx)
    else:
        ax.plot(f, 10 * np.log10(Pxx))

    ax.grid()
    plt.ylabel(ylabel)
    plt.xlabel(xlabel)
    plt.title(title)

    return ax, f, Pxx


def float_to_IQ(data):
    """
    Converts signal from floats to complex to get IQ pairs. (Hilbert transform)

    Parameters
    ----------
    data : np.array
        Input signal.

    Returns
    -------
    np.array (complex)
        Complex array of input signals I = np.abs(x), Q = np.angle(x)

    """
    return signal.hilbert(data)


# https://github.com/scipy/scipy/blob/v0.16.0/scipy/stats/stats.py#L1963
def signaltonoise(a, axis=0, ddof=0):
    """
    Compute signal-to-noise ratio of the input data.

    Returns the signal-to-noise ratio of `a`, here defined as the mean
    divided by the standard deviation.

    Parameters
    ----------
    a : array_like
        An array_like object containing the sample data.
    axis : int or None, optional
        Axis along which to operate. Default is 0. If None, compute over
        the whole array `a`.
    ddof : int, optional
        Degrees of freedom correction for standard deviation. Default is 0.

    Returns
    -------
    s2n : ndarray
        The mean to standard deviation ratio(s) along `axis`, or 0 where the
        standard deviation is 0.

    """
    a = np.asanyarray(a)
    m = a.mean(axis)
    sd = a.std(axis=axis, ddof=ddof)
    return np.where(sd == 0, 0, m / sd)


def denoise(data, threshold=1 / 50):
    """
    Perform denoising based on zeroing frequences with PSD amplitude < threshold * max(PSD).

    Parameters
    ----------
    data : np.ndarray
        Input signal.
    threshold : TYPE, optional
        Fraction of of max PSD, bellow which everything become 0. The default is 1 / 50.

    Returns
    -------
    filtered : np.ndarray
        Filtered signal.

    """
    nperseg = data.shape[0]
    f, t, Zxx = signal.stft(data, fs=1, nperseg=nperseg)
    amp = np.max(Zxx)
    Zxx = np.where(np.abs(Zxx) >= amp * threshold, Zxx, 0)
    _, filtered = signal.istft(Zxx, 1)
    return filtered