Python Numpy Additive White Gaussian Noise Function

user_2611189 0
Viewed 1461

Continuing from this thread, I need a function that does Additive White Gaussian Noise (AWGN) on my input signal.

This is my problem:

  • unable to scale to multiple channels
  • unable to scale to multiple batch
  • scale not on individual signal level

important conditions:

  • accepts numpy array of any dimension, as long as the last axis is time
  • in numpy.random.normal, the scale or standard deviation (SD) is not global, but depend on the SD of each signal. Unless my intended implementation for AWGN is wrong, that SD should be set as the SD of the entire dataset or hardcoded?

What I've done so far:

import numpy as np
import matplotlib.pyplot as plt

def add_noise(data): # assume data shape is (batch,channel,time), but it can also support (batch,time), (batch,anything,whatever,channel,time)

    time_axis = len(data.shape)-1
    target_snr_db = 20
    
    data_watts = data ** 2
    sig_avg_watts = np.mean(data_watts, axis=time_axis)
    sig_avg_db = 10 * np.log10(sig_avg_watts)
    noise_avg_db = sig_avg_db - target_snr_db
    noise_avg_watts = 10 ** (noise_avg_db / 10)
    mean_noise = 0
    noise_volts = np.random.normal(mean_noise, np.sqrt(noise_avg_watts), data.shape) # <-- problem here
    # add noise to the original signal
    noise_data = data + noise_volts
    return noise_data

This is ok, lets say we are passing a signal with 1 channel (1,1,1000):

x = np.random.rand(1,1,1000)
plt.plot(x[0,0])
plt.show()

y = add_awgn_noise(x)
plt.plot(y[0,0])
plt.show()

This is not ok, imagine passing a signal with 10 channels (1,10,1000)

x = np.random.rand(1,10,1000)
y = add_awgn_noise(x)

This is not ok, imagine passing 10 signals with 10 channels (10,10,1000)

x = np.random.rand(1,10,1000)
y = add_awgn_noise(x)
numpy python
1 Answers

In order to allow the differently shaped arrays in the np.random.normal call to be broadcast together, you have to manually tell numpy which axis you want to broadcast noise_avg_watts along:

noise_volts = np.random.normal(mean_noise, np.sqrt(noise_avg_watts)[..., np.newaxis], data.shape)

Notice the [..., np.newaxis].

If I understand correctly, noise_avg_watts is an array shaped like (s[0], s[1], ..., s[n-1]) where the data is shaped like (s[0], s[1], ..., s[n-1], s[n]); i.e. it has one scalar number per time series (and the last axis is the time series axis). So, in order to make it compatible with data's shape, we want to "broadcast" the average value across the time series axis, which is what we're doing with [..., np.newaxis].

This makes numpy treat this array as if it was the same shape as data, with all of the values in the last axis being the same (since it's an average for that axis) for each (n-1)-tuple of indices for the previous (n-1) axes. In other words, for each (signal, channel) pair, the noise will have the same standard deviation across the whole time series of the signal.

user_61174260
Related