Data augmentation layer based on physical model for time series data

Question

I am quite new to the Keras API, so forgive me if I use incorrect terminology and for my lack of knowledge about the API.

This is for a mathematical (wave modelling) research project and I am quite open to any suggestions and will be happy to modify my current implementation if it means I am using the API more effectively.

Summary of what I want:

I would like to perform data augmentation on my input layer before it is passed onto the first hidden layer. I want this to be done in a way where I am randomly generating 3 values (1 scalar, 2 vector) on each epoch which I pass into a mathematical function that creates a noise vector the same length as the input, that I will later add to the input layer.

Basically, I want to reduce overfitting on our synthetically generated training data.

I have seen that you can use Python classes to design your own layers and so I wanted to use them as a way of modularising my code, and potentially gaining performance, if it is even possible.

What I require:

Based on an input size of 2501 measurements in time (i.e., #features in x_train), create 3 variables:

• theta_w - scalar - Uniform(0, 2pi)
• r1, r2 - vector [#features by 1] - Normal(0, 1)

where I will take theta_w, r1, and r2, and substitute them into several mathematical functions to then end up with a (2501 x 1) tensor, which I can then add onto the input layer to create a "noised input layer".

At the moment this is what I have:

from keras import layers
import tensorflow as tf
import numpy as np

class WaveNoise(layers.Layer):
    def __init__(self, U10, F, **kwargs):
        super(WaveNoise, self).__init__(**kwargs)
        self.U10 = U10
        self.F = F
        
        # Initialise values that will be used in formulas
        self.alpha = 0.076 * (self.U10**2 / (self.F * self.GRAVITY)) ** 0.22
        ...

    def build(self, input_shape):
        super(WaveNoise, self).build(input_shape)
        # These only depend on the input shape, so they only need to be initialised
        self.time = tf.linspace([0.0], [self.SAMPLING_PERIOD * input_shape[1]], input_shape[1], name="time", axis=1)
        ...

        # These are the values which I want to update on each epoch
        # I believe self.build() is only called once and I should be placing this in self.call(), but when I do that, I get a ValueError for the initial_value arg
        # Also should these be non-trainable?
        theta_w_init = tf.random_uniform_initializer(minval=0.0, maxval=2 * np.pi)
        self.theta_w = tf.Variable(initial_value=theta_w_init(shape=(1, 1)), trainable=True)

        r_init = tf.random_normal_initializer(mean=0.0, stddev=1.0)
        self.r1 = tf.Variable(initial_value=r_init(shape=input_shape[1:]), trainable=True)
        self.r2 = tf.Variable(initial_value=r_init(shape=input_shape[1:]), trainable=True)

    def call(self, inputs):
        self.a = tf.math.sqrt(2 * self.E(self.omega, self.theta) * self.d_omega * self.d_theta)

        R = tf.math.sqrt(self.r1 ** 2 + self.r2 ** 2)
        Phi = self.phi((self.r1, self.r2))

        return self.sum_of_waves((R, Phi, inputs.shape))

    ... # Other methods to perform the mathematical calculations above, such as `self.E()`

I add this into my network through:

input_layer = layers.Input((num_features, 1), name="input-layer")

U10 = 15
F = 100000
wave_noise = WaveNoise(U10, F, name="wave-noise")(input_layer)

wave_noise_model = keras.Model(inputs=input_layer, outputs=wave_noise, name="noise-model")

# Add noise to input
noised_input_layer = layers.Add(name="noised-input")([input_layer, wave_noise])

I would like to know if I can improve this code in any way and if this code best captures what I want. So, if anyone has any suggestions, I would love to hear them.

Thanks for the help!