Slight variation of Tme2Vector implementation.

time2vector.py

# -*- coding: utf-8 -*-
"""
Created by Pierluigi on 2020-09-30
"""
from tensorflow.keras.layers import Layer
from tensorflow.keras import backend as K


class Time2Vector(Layer):
    """ Time2Vector

    This is a slight variation of the layer suggested here: https://arxiv.org/abs/1907.05321

    This layer expect as input a tensor like (batch, time_steps, signals) to each
    signal it applies the augmentation described in the paper. Note that this
    layer will, potentially, return a tensor bigger than the one in input, in
    particular if you have `num_signals` as input, you will obtain `(1 + sins) * num_signals`
    as output.

    :param sins: number of different sinusoids to be used
    :param decompose: if True will apply additive timeseries decomposition
    """
    def __init__(self, sins=2, decompose=False, **kwargs):
        super(Time2Vector, self).__init__(**kwargs)
        self.sins = sins
        self.decompose = decompose

    def build(self, input_shape):
        assert len(input_shape) == 3, f'input tensor should be 3 dimensions: (batch, time_steps, signals) got {input_shape}'
        _, self.time_steps, self.signals = input_shape
        self.w = self.add_weight(name='weight',
                                 shape=(self.time_steps, self.signals, ),
                                 initializer='uniform',
                                 trainable=True)

        self.b = self.add_weight(name='bias',
                                 shape=(self.signals, ),
                                 initializer='uniform',
                                 trainable=True)
        self.freqs = []
        self.phases = []
        for i in range(self.sins):
            self.freqs.append(self.add_weight(name=f'freq_{i:03d}',
                                              shape=(self.time_steps, self.signals, ),
                                              initializer='uniform',
                                              trainable=True))

            self.phases.append(self.add_weight(name=f'phase_{i:03d}',
                                               shape=(self.signals, ),
                                               initializer='uniform',
                                               trainable=True))

    def call(self, x, **kwargs):
        res = [self.w * x + self.b]
        if self.decompose:
            x -= res[-1]
        for i in range(self.sins):
            # added a scale factor (i + 1) / (self.sins + 1) in order to
            # be sure that each sins differs from the other
            res.append(K.sin((i + 1) / self.sins * x * self.freqs[i] + self.phases[i]))
            if self.decompose:
                x -= res[-1]
        return K.concatenate(res, axis=-1)

    def compute_output_shape(self, input_shape):
        return self.time_steps, (1 + self.sins) * self.signals

    def get_config(self):
        config = super().get_config().copy()
        config.update({
            'sins': self.sins,
            'decompose': self.decompose,
        })
        return config