bootstrap_snr.py

import numpy as np
import matplotlib.pyplot as plt
plt.ion()

rand = np.random.RandomState()
rms = lambda x, axis=-1: np.sqrt((x ** 2).sum(axis=axis))
ms = lambda x, axis=-1: (x ** 2).sum(axis=axis)


# bootstrap to guess SNR
def bootstrap_snr(ep, n_bootstrap=1000, n_mean=4000, do_plot=True):
    x = ep.get_data()
    n_mean = 2 * (n_mean // 2)
    n_ep, n_ch, _ = x.shape
    sig_noise_lev = np.zeros((n_bootstrap, n_ch))
    noise_lev = np.zeros((n_bootstrap, n_ch))
    sign = np.concatenate((-np.ones(n_mean / 2),
                           np.ones(n_mean / 2)))[:, np.newaxis, np.newaxis]
    for bi in range(n_bootstrap):
        sig_noise_lev[bi] = ms(x[rand.randint(0, n_ep, n_mean)].mean(0))
        noise_lev[bi] = ms((x[rand.randint(0, n_ep, n_mean)]
                            * sign).mean(0))

    snr = 10 * (np.log10((sig_noise_lev.mean(axis=0) - noise_lev.mean(axis=0)))
                - np.log10(noise_lev * n_mean))
    snr_mean = snr.mean(axis=0)
    snr_std = snr.std(axis=0)

    if do_plot:
        plt.errorbar(range(n_ch), snr_mean, snr_std, fmt=None, ecolor='k')
        plt.plot(range(n_ch), snr_mean, 'o')
        x_lim = [-0.5, n_ch - 0.5]
        plt.plot(x_lim, -10 * np.log10([n_ep] * 2) + 6, 'k:')
        plt.plot(x_lim, -10 * np.log10([n_ep] * 2), 'k-')
        plt.xlim(x_lim)
        plt.xticks(range(n_ch), ep.ch_names)
        plt.ylabel('SNR (dB)')

    return snr_mean, snr_std


# test it out -- should also do this as a function of time
if __name__ == '__main__':
    sig_len = 100
    snr_db = -10
    snr = 10 ** (snr_db / 20.)
    sig = np.sin(np.arange(sig_len) /
                 float(sig_len) * 2 * np.pi) / np.sqrt(2) * 2
    import mne
    ep_info = mne.create_info(['Ch1'], sfreq=sig_len)
    n_ep = 1000
    ep_data = (sig[np.newaxis, np.newaxis, :] +
               rand.randn(n_ep, 1, sig_len) / snr)
    ep_events = np.concatenate([np.atleast_2d(np.arange(n_ep) * sig_len)] +
                               2 * [np.atleast_2d(np.ones(n_ep))]).T
    ep = mne.EpochsArray(ep_data, ep_info, ep_events)
    bootstrap_snr(ep, n_bootstrap=1000, n_mean=2000)