Source code for node_wavelets

# This file is part of Sympathy for Data.
# Copyright (c) 2018, Combine Control Systems AB
#
# SYMPATHY FOR DATA COMMERCIAL LICENSE
# You should have received a link to the License with Sympathy for Data.
import numpy as np
import pywt

from sympathy.api import node
from sympathy.api.nodeconfig import Port, Ports, Tag, Tags


wavelets = {
    'ricker': 'mexh',
    'morlet (real)': 'morl',
}




class ContinuousWaveletTransform(node.Node):
    """
    Continuous Wavelet Transform (CWT) is similar to
    :ref:`Short Time Fourier Transform<com.sympathyfordata.timeseriesanalysis.spectral_transform>`
    in that it reveals the frequency contents of a signal, and also shows how
    that frequency content changes over time.
    """
    name = 'Continuous wavelet transform'
    author = 'Mathias Broxvall'
    icon = 'fourier.svg'
    description = (
        'Decompose a signal into multiple different scales/frequencies and show '
        'how that frequency component of the signal varies over time.'
    )
    nodeid = 'com.sympathyfordata.timeseriesanalysis.continous_wavelet_transform'
    tags = Tags(Tag.Analysis.SignalProcessing)
    related = [
        'com.sympathyfordata.timeseriesanalysis.frequency_transform',
        'com.sympathyfordata.timeseriesanalysis.spectral_transform',
    ]

    parameters = node.parameters()
    parameters.set_string(
        'wavelet',
        label='Wavelet',
        value='ricker',
        description='Selects wavelet to apply',
        editor=node.editors.combo_editor(
            options=list(wavelets.keys())))
    parameters.set_integer(
        'scales',
        label='Number of scales',
        value=32,
        description='Number of scale values to generate, always start at '
                    'scale 1')
    parameters.set_boolean(
        'rescale',
        label='Rescale output',
        value=False,
        description='Multiplies each scale output with 1/scale as a '
                    'normalization')

    inputs = Ports([
        Port.Table('time domain', name='time domain')
    ])
    outputs = Ports([
        Port.Table('wavelet domain', name='wavelet domain')
    ])

    def execute(self, node_context):
        time_tbl = node_context.input['time domain']
        out_tbl = node_context.output['wavelet domain']
        scales = node_context.parameters['scales'].value
        rescale = node_context.parameters['rescale'].value

        wavelet = wavelets[node_context.parameters['wavelet'].value]

        for col in time_tbl.cols():
            name = col.name
            widths = np.array(range(1, scales+1))
            coefs, _freqs = pywt.cwt(col.data, widths, wavelet)
            for i, c in enumerate(coefs):
                if rescale:
                    c = c / widths[i]
                out_tbl[f'{name}@{i+1}'] = c