Python while loop mulitple plotting

时间:2015-06-15 14:34:51

标签: python matplotlib plot while-loop

maybe someone did something similiar.

I have a list of parameters:

i = 120
a = 122

n_epochs = 120
sfreq, duration = 1000., 1000
times = np.arange(0, duration, 1 / sfreq)
amp , amp2 , nse_amp = 1., 1., 0.5
nfft = 512

nse1 = np.random.rand(times.size) * nse_amp
nse2 = np.random.rand(times.size) * nse_amp
x =  amp * np.sin(2 * np.pi * 200 * times) + nse1
y =  amp * np.sin(2 * np.pi * 200 * times + np.pi/5) + nse2

n_freqs = nfft/2 + 1

and a bunch of functions (working correctly), that are needed to compute my last function.

This function is depending on multiple parameters. I want to run my function in a loop as long as i reaches a and I want to plot all the results.

My function:

def my_con(x, y, n_epochs, nfft, sfreq, con_name='coh'):
    '''Computes connectivity measure mentioned on provided signal pair and its surrogates.'''

    freqs, Rxy, Rxy_mean, Rxx_mean, Ryy_mean = compute_mean_psd_csd(x, y, n_epochs, nfft, sfreq)

    # compute surrogates
    x_surr = x.copy()
    y_surr = y.copy()
    np.random.shuffle(x_surr)
    np.random.shuffle(y_surr)
    freqs_surro, Rxy_s, Rxy_s_mean, Rxx_s_mean, Ryy_s_mean = compute_mean_psd_csd(x_surr, y_surr, n_epochs, nfft, sfreq)

    if con_name == 'coh':
        coh = my_coherence(n_freqs, Rxy_mean, Rxx_mean, Ryy_mean)
        coh_surro = my_coherence(n_freqs, Rxy_s_mean, Rxx_s_mean, Ryy_s_mean)
        return coh, coh_surro, freqs, freqs_surro

Now the task to do: I want this function to work in a loop. While or for I worked on both but it didn't work out.

con, con_surro, freqs, freqs_surro = my_con(x, y, n_epochs, nfft, sfreq, con_name) # for doing every con_surro calc. too add [i]

Also my plotting set up:

pl.figure('Connectivity')
pl.plot(freqs, con)
pl.plot(freqs_surro, con_surro)
pl.legend(['Con', 'Surrogates'])
pl.tight_layout()
pl.show(True)

Anyone who is capable of helping me ?!

Thanks in advance !

1 个答案:

答案 0 :(得分:0)

import sys
import numpy as np
import matplotlib.pyplot as pl
import matplotlib.mlab as mlab

n_epochs = 120
a = 140
while n_epochs <= a:
    sfreq = 1000.
    duration = 1000.
    times = np.arange(0, duration, 1 / sfreq)
    amp = 1.
    amp2 = 1.
    nse_amp = 0.5
    nfft = 512
    nse1 = np.random.rand(times.size) * nse_amp
    nse2 = np.random.rand(times.size) * nse_amp
    x =  amp * np.sin(2 * np.pi * 200 * times) + nse1
    y =  amp * np.sin(2 * np.pi * 200 * times + np.pi/5) + nse2

    shift = 100  # integer
    assert shift < sfreq * duration, 'Choose a smaller shift.'
    #y = amp2 * np.roll(x, shift) + nse2

    # coherence using mlab function
    cohxy, freqs = mlab.cohere(x, y, Fs=sfreq, NFFT=nfft)

    n_freqs = nfft/2 + 1

    def compute_mean_psd_csd(x, y, n_epochs, nfft, sfreq):
        '''Computes mean of PSD and CSD for signals.'''
        x2 = np.array_split(x, n_epochs)
        y2 = np.array_split(y, n_epochs)

        Rxy = np.zeros((n_epochs, n_freqs), dtype=np.complex)
        Rxx = np.zeros((n_epochs, n_freqs), dtype=np.complex)
        Ryy = np.zeros((n_epochs, n_freqs), dtype=np.complex)

        for i in range(n_epochs):
            Rxy[i], freqs = mlab.csd(x2[i], y2[i], NFFT=nfft, Fs=sfreq)
            Rxx[i], _ = mlab.psd(x2[i], NFFT=nfft, Fs=sfreq)
            Ryy[i], _ = mlab.psd(y2[i], NFFT=nfft, Fs=sfreq)

        Rxy_mean = np.mean(Rxy, axis=0)
        Rxx_mean = np.mean(Rxx, axis=0)
        Ryy_mean = np.mean(Ryy, axis=0)

        return freqs, Rxy, Rxy_mean, np.real(Rxx_mean), np.real(Ryy_mean)

    def my_coherence(n_freqs, Rxy_mean, Rxx_mean, Ryy_mean):
        ''' Computes coherence. '''

        coh = np.zeros((n_freqs))
        for i in range(0, n_freqs):
            coh[i] = np.abs(Rxy_mean[i]) / np.sqrt(Rxx_mean[i] * Ryy_mean[i])

        return coh

    def my_con(x, y, n_epochs, nfft, sfreq, con_name='coh'):
        '''Computes connectivity measure mentioned on provided signal pair and its surrogates.'''

        freqs, Rxy, Rxy_mean, Rxx_mean, Ryy_mean = compute_mean_psd_csd(x, y, n_epochs, nfft, sfreq)

        # compute surrogates
        x_surr = x.copy()
        y_surr = y.copy()
        np.random.shuffle(x_surr)
        np.random.shuffle(y_surr)
        freqs_surro, Rxy_s, Rxy_s_mean, Rxx_s_mean, Ryy_s_mean = compute_mean_psd_csd(x_surr, y_surr, n_epochs, nfft, sfreq)

        if con_name == 'coh':
            coh = my_coherence(n_freqs, Rxy_mean, Rxx_mean, Ryy_mean)
            coh_surro = my_coherence(n_freqs, Rxy_s_mean, Rxx_s_mean, Ryy_s_mean)
            return coh, coh_surro, freqs, freqs_surro

        if con_name == '':
            print 'Please provided the connectivity method to use.'
            sys.exit()
        else:
            print 'Connectivity method unrecognized.'
            sys.exit()


    con_name = 'coh'


    con, con_surro, freqs, freqs_surro = my_con(x, y, n_epochs, nfft, sfreq, con_name) # for doing every con_surro calc. too add [i]

    # coherence using mlab function
    #cohxy, freqs = mlab.cohere(x, y, Fs=sfreq, NFFT=nfft)
    #pl.plot(freqs, cohxy)

    # plot results
    pl.figure('Connectivity')
    pl.plot(freqs, con)
    pl.plot(freqs_surro, con_surro)
    pl.legend(['Con', 'Surrogates'])
    pl.tight_layout()
    pl.show(False)
    pl.draw()
    n_epochs = n_epochs +1

that might solve my own problem. Someone wants to give some suggestion ?!