Source code :: main

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import matplotlib 

#matplotlib.rcParams['backend'] = 'TkAgg'
#matplotlib.rcParams['backend'] = 'WXAgg'
#matplotlib.rcParams['backend'] = 'GTKAgg'
#matplotlib.rcParams['backend'] = 'QtAgg'
##matplotlib.rcParams['backend'] = 'Qt4Agg'
matplotlib.rcParams['backend'] = 'Agg'


#TODO vypočátávat raději crosscoherenci? 

from matplotlib.pyplot import *
from numpy import *
import os
from scipy import signal
from scipy.signal import fftconvolve
from multiprocessing import Process, Pool, cpu_count
import time
from scipy.stats.mstats  import mquantiles
import matplotlib.animation as manimation
import re
from CWT import cwt
from fftshift import fftshift
from SoundGenerator import soundGenerator
from numexpr import evaluate

global mirnov, ring
mirnov = False
ring=True
animation = False

RingCoilOrientation = {12000:(-1,1,1,1,-1,1,-1,-1,-1,1,-1,1,-1,-1,-1,-1),
 0:(-1,-1,1,1,-1,1,-1,1,-1,1,-1,-1,-1,-1,-1,-1),
 18000:(+1,+1,-1,-1,+1,-1,+1,-1,+1,-1,+1,+1,+1,+1,+1,-1) }

MirnovCoilOrientation = array((-1,-1,1,1))
MirnovEffectiveArea = array((1,1,1,1))


EffectiveArea = {12000:[ 0.00209134 , 0.00446023 , 0.00428289  ,0.00374073 , 0.00320856 , 0.00182771,
  0.00267794,  0.00284156 , 0.00145525  ,0.00222005  ,0.00206269 , 0.00108452,
  0.00054154 , 0.00078395 , 0.00186135 , 0.00274475],
  18000:[ 0.00209134 , 0.00446023 , 0.00428289  ,0.00374073 , 0.00320856 , 0.00182771,
  0.00267794,  0.00284156 , 0.00145525  ,0.00222005  ,0.00206269 , 0.00108452,
  0.00054154 , 0.00078395 , 0.00186135 , 0.00274475],
  1:[ 0.00209134 , 0.00446023 , 0.00428289  ,0.00374073 , 0.00320856 , 0.00252771,
  0.00267794,  0.00284156 , 0.00145525  ,0.00222005  ,0.00206269 , 0.00005452,
  0.00254154 , 0.00078395 , 0.00186135 , 0.00274475],
   0:[    2.45766087e-03  , 3.33461911e-03 ,  3.21871869e-03 ,  4.5e-03,
   2.39870295e-03  , 3.99826880e-03  , 3.47274534e-03 ,  3.93465853e-03,
   2.25236901e-03  , 3.22086160e-03 ,  3.27836414e-03 ,  6.87539836e-05,
   3.77966762e-03  , 1.78949877e-03  , 3.19212003e-03  , 3.50243652e-03]}



def smooth(x,window_len=11,window='hanning',axis=-1):
    if not x.ndim in (1,2):
            raise ValueError, "smooth only accepts 1 or 2 dimension arrays."
    if x.shape[axis] < window_len:
            raise ValueError, "Input vector needs to be bigger than window size."
    if window_len < 3:
            return x
        
    x = atleast_2d(x)
    x = swapaxes(x, -1, axis)
    y = empty_like(x)
    
    
    if isinstance(window, str) or type(window) is tuple:
        w = signal.get_window(window, window_len)
    else:
        w = np.asarray(window)
        if len(w.shape) != 1:
            raise ValueError('window must be 1-D')
        if  w.shape[0] > x.shape[-1]:
            raise ValueError('window is longer than x.')
        window_len = w.shape[0]
    
    
    for i in range(x.shape[0]):
        s=r_[2*x[i,0]-x[i,window_len-1::-1],x[i,:],2*x[i,-1]-x[i,-1:-window_len:-1]]
        y[i,:]=fftconvolve(s,w/w.sum(),mode='same')[window_len:-window_len+1]     
    
    y = swapaxes(y, -1, axis)

    return squeeze(y)



class LogFormatterTeXExponent(LogFormatter, object):
    """Extends pylab.LogFormatter to use 
    tex notation for tick labels."""
    
    def __init__(self, *args, **kwargs):
        super(LogFormatterTeXExponent, 
              self).__init__(*args, **kwargs)
        
    def __call__(self, *args, **kwargs):
        """Wrap call to parent class with 
        change to tex notation."""
        label = super(LogFormatterTeXExponent, 
                      self).__call__(*args, **kwargs)
        label = label.replace('-','')
        
        x = float(label)
        if abs(log10(x)) >2:                      
            label = re.sub(r'e(\S)0?(\d+)',r'\\cdot 10^{\1\2}',str(label))
            label = "$" + label + "$"
        else:
            n = max(0,-int(log10(x)-1))
            label = ('$%.'+str(n)+'f$')%x

        return label
        


def mad(x,axis = -1):
    #Median absolute deviation
    m = np.median(x,axis=axis)
    x = np.swapaxes(x, 0, axis)
    s = np.median(abs(x-m),axis=0)*1.4826
    return s

def DFT(x, y,nf):
    
    Ftheta = arange(-nf/2,nf/2)

    E = exp(-1j*outer(Ftheta,x))
    #print y.shape, E.shape, nf, linalg.pinv(E).shape,Ftheta.shape, x.shape
    Fx_ = dot(y,conj(linalg.pinv(E)))*nf
    
    return  Fx_




def LoadSignal2():
    from pygolem_lite import Shot
    from pygolem_lite.modules import list2array,deconvolveExp,save_adv,saveconst
    global mirnov, ring

    #print ring, mirnov
    Data = Shot()
    gd = Shot().get_data
    Ring,Mirnov = None,None
    #gd('any', return_channel = True)
    #tvec, data  = gd('NIoctopus', return_channel = False)
    #print das
    #print data

    #exit()

    try:
        if ring:
            tvec, Ring  = gd('NIoctopus', return_channel = False)
            #Ring = list2array( Data[das, range(16)] )
            #tvec,Ring = Ring[:,0],Ring[:,1:]

    except Exception as e:
        ring = False
        print 'No data from ring', e
    #if mirnov:
        #das, m1  = gd('any', 'mirnov_1', return_channel = True)
        #das, m5  = gd('any', 'mirnov_5', return_channel = True)
        #das, m9  = gd('any', 'mirnov_9', return_channel = True)
        #das, m13 = gd('any', 'mirnov_13', return_channel = True)
        #Mirnov = list2array( Data[das, [m1, m5, m9, m13]] )
    ##tve
    
    #exit()
    try:
        if mirnov:
            das, m1  = gd('any', 'mirnov_1', return_channel = True)
            das, m5  = gd('any', 'mirnov_5', return_channel = True)
            das, m9  = gd('any', 'mirnov_9', return_channel = True)
            das, m13 = gd('any', 'mirnov_13', return_channel = True)
            Mirnov = list2array( Data[das, [m1, m5, m9, m13]] )
            tvec,Mirnov = Mirnov[:,0],Mirnov[:,1:]
            Mirnov/= 0.0005

    except:
        mirnov = False
        print 'no data from mirnov'
    shot  = Data.shot_num
    
    if Ring!= None and Mirnov!= None:
        data = c_[Ring, Mirnov]
    elif Ring!= None:
        data = Ring
    else:
        data = Mirnov

    plasma = Data['plasma']

    tstart = Data['plasma_start']
    tend = Data['plasma_end']
    
    tvec, data = PrepareSignal(tvec, data,tstart,tend,shot)
    return tvec, data 


def LoadSignal(shot):
    print 'loading'
    from golem_data import golem_data


    obj = golem_data(shot, 'niturbo')
    data = obj.data
    tvec = obj.tvec
    tend   = golem_data(shot, 'plasma_end').data
    tstart = golem_data(shot, 'plasma_start').data

    data2 = golem_data(shot, 'nistandard6132').data
    data = c_[data, data2]/r_[ones(16),ones(4)*0.0005][None,:]
    
    tvec, data = PrepareSignal(tvec, data,tstart,tend,shot)
    return tvec, data 
    
        
def PrepareSignal(tvec, data,tstart, tfinish,shot ):
    print 'preparing'
    if shot < 12000:
        calib_shot = 0 

    elif shot < 18000:
        calib_shot = 12000
    else:
        calib_shot = 18000
    
    if shot > 12000:
        if mirnov and ring:
            AEff = r_[array(EffectiveArea[calib_shot]),MirnovEffectiveArea]
            polarity = r_[array(RingCoilOrientation[calib_shot]),MirnovCoilOrientation]
        elif mirnov:
            AEff = MirnovEffectiveArea
            polarity = MirnovCoilOrientation
        else:
            AEff = array(EffectiveArea[calib_shot])
            polarity = array(RingCoilOrientation[calib_shot])
  
    data/= (AEff*polarity)[None,:]
    data = signal.detrend(data, axis=0)

        
    data = cumsum(data, axis=0)*1e-6

    N = 101
    taps = signal.firwin(N, 1e3/(0.5*1e6) , window='hamming', pass_zero=True)   
    data -= smooth(data,window_len=101,window=taps,axis=0)

    nstart = argmin(abs(tvec-tstart))
    nend = argmin(abs(tvec-tfinish))
    
    data = data[nstart:nend,:]
    data = signal.detrend(data, axis=0)

    tvec = tvec[nstart:nend]

    return tvec, data


def GenerateSound(mode_signal,modes):

  
    #use only one "polarization" of the signal
    mode_signal = real(mode_signal)
    
    path = './mp3/'
    if not os.path.exists(path): os.makedirs(path)
    
    
    t2= time.time()
    
    pool = Pool(cpu_count())
    out = pool.map(soundGenerator,[(mode_signal[:,i],2000, path+'sound'+str(m)) for i,m in modes])
    pool.close()
    pool.join()


    print 'generated sound', time.time()-t2


def PlotSpecrograms(freq, field, t, modes, logscale=True):
    print 'PlotSpecrograms'
    
    field = log(1+(field/std(field))**2)/2
    
    vmin = amin(field)
    vmax = amax(field)
    
    fig = figure('spectrogram')

    ax = fig.add_axes([0.1, 0.1, 0.8, 0.85])
    if logscale:
        ax.set_yscale('log', nonposy='clip')
        ax.yaxis.set_minor_formatter(LogFormatterTeXExponent(base=10,
            labelOnlyBase=False))        
        ax.yaxis.set_major_formatter(LogFormatterTeXExponent(base=10,
            labelOnlyBase=False))
            
    img = ax.imshow(zeros((1,1)), extent=[t[0]*1e3,t[-1]*1e3 ,freq[-1]/1e3, freq[0]/1e3], aspect='auto',vmin=vmin, vmax=vmax,interpolation='bicubic') 
    minorLocator   = MultipleLocator(1)

    ax.xaxis.set_minor_locator(minorLocator)
    ax.axis([t[0]*1e3,t[-1]*1e3 ,amin(freq)/1e3, amax(freq)/1e3])
    ax.set_xlabel('time [ms]')
    ax.set_ylabel('Frequency [kHz]')

    t = time.time()
    for i,m in modes:
        #print 'i:%d'%i
        img.set_data((field[...,size(field,2)-i-1]))
        fig.savefig('./graphs/spectrogram'+str(m)+'.png')
        
    print 'plotted spectograms', time.time()-t

    fig.clf()
    

def AnalyzeSpectrograms(spectrogram,tvec,fvec,fmin,fmax,log_scale=True):
    print 'AnalyzeSpectrograms'
    #max_freq = fvec[argmax(spectrogram, axis=0)]
    #print spectrogram.shape
    #exit()
    n_modes = spectrogram.shape[2]
    ind_freq = argmax(spectrogram, axis=0)

    ind_freq = int_(signal.medfilt(ind_freq, (31,1)))
    max_freq = fvec[ind_freq]

    max_amplitude = amax(spectrogram, axis=0)
    max_amplitude /= mean(max_amplitude)
 
    fig = figure(figsize=(8,8))
    subplots_adjust(hspace=0.05, wspace = 0)
    modes = [-2, -1, 0, 1, 2]

    ax = fig.add_subplot(211)
    tvec_spec = linspace(tvec[0],tvec[-1], size(spectrogram,1))*1e3
    #print range(min(5, n_modes-1),0,-1), max_amplitude.shape
    plots = [ax.plot(tvec_spec,max_amplitude[:,i]*10)[0] for i in modes]
    leg = ax.legend(plots, ['M=%d'%i for i in modes], loc='best', fancybox=True)
    leg.get_frame().set_alpha(0.5)
    ax.set_ylabel('Amplitude [a.u.]')
    ax.xaxis.set_major_formatter(NullFormatter() )
    ax.set_xlim(tvec[0]*1e3,tvec[-1]*1e3)
    
    for label in leg.get_lines():
        label.set_linewidth(3)
    ax = fig.add_subplot(212)
    
    if log_scale:
        ax.set_yscale('log', nonposy='clip')        
        ax.yaxis.set_minor_formatter(LogFormatterTeXExponent(base=10,
            labelOnlyBase=False))    
        ax.yaxis.set_major_formatter(LogFormatterTeXExponent(base=10,
            labelOnlyBase=False))  
            
    #modes = range(min(5, n_modes-1),0,-1)
    #print modes
    #exit()
    plots = [ax.plot(tvec_spec,max_freq[:,i]/1e3,linewidth=1)[0] for i in modes]
    
    
    leg = ax.legend(plots, ['M=%d'%i for i in modes], loc='best', fancybox=True)
    leg.get_frame().set_alpha(0.5)
    
    for label in leg.get_lines():
        label.set_linewidth(3)

    ax.set_ylim((fmin-1)/1e3,(fmax+1)/1e3)

    ax.set_xlabel('t [ms]')
    ax.set_ylabel('Frequency [kHz]')
    ax.set_xlim(tvec[0]*1e3,tvec[-1]*1e3)
    
    fig.savefig('graphs/analyze.png')#,bbox_inches='tight'
    close()
   


def  AnimateSpectrogram(x,spectrogram,fvec,tvec,theta,nmods,
            nch,framesamp,hopsamp,log_scale=True):

    if not animation:
        return 
    #plot animation 
    nt = size(x,0)

    
    fig = figure(figsize=(10,5))
    ax1 = fig.add_subplot(121)
    ax2 = fig.add_subplot(122)
    

    mag_img = ax1.imshow(zeros((framesamp,nmods/2+1 )), 
        extent=[fvec[0]/1e3,fvec[-1]/1e3,-nmods/2+0.5,nmods/2+0.5],aspect='auto', 
            animated=True,interpolation='nearest')

    mag_img.set_cmap('YlOrBr')
  
    plot2, = ax2.plot([],[], linewidth=0.3, animated=False)

    ax2.set_ylim(-0.3,2*pi+0.1)
    ax2.set_ylabel('poloidal angle [rad]')

    ax2.axvline(x=0)

    Dt = tvec[framesamp]-tvec[0]
    ax2.set_xlim(-Dt/2*1e6, Dt/2*1e6)
    ax2.set_xlabel('$\Delta$t [$\mu$s]')
    ax1.set_xlabel('f [kHz]')

    time_text = ax2.text(0.05, 0.9, '', transform=ax2.transAxes,animated=False)
    
    ax1.yaxis.set_major_locator(MultipleLocator(1))
    ax2.xaxis.grid(color='r', linestyle='-', linewidth=.2)
    ax1.set_ylabel('Poloidal mode number M')
    
    if log_scale:
        ax1.set_xscale('log', nonposy='clip')
        ax_format = LogFormatterTeXExponent(base=10,labelOnlyBase=False)
        ax1.xaxis.set_minor_formatter(ax_format)        
        ax1.xaxis.set_major_formatter(ax_format)    
            
    
    def init():
        plot2.set_data([],[])
        return plot2, 

        
    t = 0
    
    FFMpegWriter = manimation.writers['mencoder']
    metadata = dict(title='NTM', artist='Matplotlib', comment='Animated CWT spectrogram')
    writer = FFMpegWriter(fps=10, metadata=metadata,codec='mpeg4',bitrate=2000)

    with writer.saving(fig, 'graphs/AnimCWT.mp4', 100):
        for j in xrange(0, size(x,0)-framesamp, hopsamp):

            i_spec = ((j+framesamp/2)*size(spectrogram,1))/nt
            
            time_text.set_text('i=%d t=%.2fms'%(j+framesamp/2,1e3*tvec[j+framesamp/2]))

            y = copy(x[j:j+framesamp, : ])
            y/= mad(y,axis=0)[None,:]


            image = spectrogram[:,i_spec,:].T

            image = hypot(image/std(image), 1)
            
            mag_img.set_data(image)

            mag_img.set_clim(1,amax(image[nmods/2+2:,:]))
            T = tile(tvec[:framesamp ]-tvec[:framesamp].mean(),(1,nch)).T
            T[::framesamp] = nan
            
            plot2.set_data(T*1e6, (y/8+theta[None,:]).T.ravel())
            
            fps = 1/(time.time()-t)
            t = time.time()

            sys.stdout.write('\r  plotting %2.1f%%  fps:%2.1f'%(j*100./nt,fps))
            sys.stdout.flush()
            writer.grab_frame() 

      
    
    

def  multi_cwt(x, fs, framesz, hop,tvec=None):
    
    
    framesamp = int(framesz*fs)
    hopsamp = int(hop*fs)
    if tvec == None:
        tvec = arange(size(x,0))/fs
        
        
    nt, nch =  shape(x)
    nmods = min(16, nch)

    fmax = 8e4
    fmin = 3e3
    dj = 0.05  #frequency resolution of CWT
    p = 10    #length of the wavelet
    
    #calculate intensity envelope ma moving STD
    filtered_std = sqrt(abs(smooth(x**2,window_len=1e3,window='hanning',axis=0)))

    x /= filtered_std/mean(filtered_std,axis=1)[:,None]

   
    #poloidal angles of the coils in the !magnetics! coordinates
    global ring, mirnov
    if ring and mirnov:
        theta = load('phi.npy')
        theta = median(theta,axis=0)
        theta += 2*pi*r_[arange(16)/16.,arange(4)/4.]
    elif mirnov:
        theta = 2*pi*arange(4)/4.
    elif ring:
        theta = 2*pi*arange(16)/16.

        
  
    #DFT over the channels with nonequidistant sampling
    fx = DFT(theta, x,nmods)
    
    print 'calc spectrograms'
  

    pool = Pool(cpu_count())
    out = pool.map(cwt,[(fx[:,i], 1/fs, dj, p ,nt/hopsamp, fmin,fmax) for i in range(nmods)])
    pool.close()
    pool.join()
    freq = out[0][2][::-1]
    
    spectrogram = [abs(out.pop()[0]) for i in range(nmods)]
    spectrogram = dstack(spectrogram)
    
    modes= zip(arange(nmods), arange(-nmods/2, nmods/2))
    
    #create sound of the modes
    GenerateSound(fx,modes)
    
    #analyze the spectrograms
    AnalyzeSpectrograms(spectrogram[::-1,...],tvec,freq,fmin,fmax )
    spectrogram = spectrogram[::-1,:,::-1]

    #plot spectrograms
    PlotSpecrograms(freq,spectrogram,tvec,modes)
  
    #plot animated 3d spectrograms
    AnimateSpectrogram(x,spectrogram,freq,tvec,theta,nmods,nch,framesamp,hopsamp)
 
    
def calc_stft((y,w,nfmax,theta)):
    y = y/mad(y,axis=0)[None,:]
    #print y.shape
    #exit()
    fy = DFT(theta, y,16)
    f = fft.fft(fy*w[:,None], axis=0)
    #BUG RFFT from scipy.fftpack for singles is faster!!
    #print f.shape, nfmax

    return f[:nfmax,:]         
        
        
        
def stft(x, fs, framesz, hop,tvec=None):
    print 'stft'
    nmods = 16
    fmax = 8e4
    nt, nch =  x.shape

    framesamp = int(2**np.ceil(np.log2(framesz*fs)))
    hopsamp = int(hop*fs)
    w = hamming(framesamp)
   
    if tvec == None:
        tvec = arange(size(x,0))/fs
 
    nfmax = int(fmax/fs*framesamp)
    theta = load('phi.npy')

    theta  = median(theta,axis=0)
    theta += 2*pi*r_[arange(16)/16.,arange(4)/4.]
            

    pool = Pool(cpu_count())
    ind = xrange(0, size(x,0)-framesamp, hopsamp)

    out = pool.map(calc_stft,[(x[i:i+framesamp,:],w,nfmax,theta) for i in ind])
    pool.close()
    pool.join()
    
    spectrogram = [abs(out.pop()) for i in ind]
    spectrogram = dstack(spectrogram)

    spectrogram = swapaxes(dstack(spectrogram),0,2)
    fvec = arange(framesamp)*fs/(2.*framesamp)
    modes= zip(arange(nmods), arange(-nmods/2, nmods/2))

    AnalyzeSpectrograms( spectrogram,tvec,fvec,0,fmax)
    PlotSpecrograms([0,fmax],spectrogram[:,::-1,:], tvec,modes, logscale=False)


def LSQfit(tvec,data):
   
    data = signal.detrend(data,axis=0)
    data/= std(data,axis=0)
 
    
    nch = size(data,1)
    theta0 = 2*pi*r_[arange(16)/16.,arange(4)/4.]
    #TODO pro každé jinou aplitudu a offset? 
    def Model(par,tvec,show_plt=False):
        A = par[3]
        f = par[0]

        phi = par[1]
        

        M =  int(par[2])
        if M != 0:
            par[4:][par[4:] <-pi/2] +=pi/2
            par[4:][par[4:] > pi/2] -=pi/2


        Dtheta =  par[4:]
        Dtheta-= median(Dtheta)

        theta = theta0[:nch]+Dtheta
        model = sin(2*pi*(tvec *f)-theta[:,None]*M+phi)*A

        if show_plt:
            plot(model.T/2+theta0[None,:nch]/2*pi,'--',linewidth=0.3)
            plot(data/2+theta0[None,:nch]/2*pi,linewidth=0.3)
        return linalg.norm((model-data.T)/std(data,axis=0)[:,None])+linalg.norm(diff(Dtheta))

    f = linspace(0,1e6/2, size(data,0)/2+1)
    F = mean(abs(fft.rfft(data*hamming(size(data,0))[:,None],axis=0)),axis=1)
    i = argmax(F)
    f_max = sum((f*F)[i-1:i+2])/sum(F[i-1:i+2])

    f_max = max(29000,f_max)
    
    theta_ = median(load('phi2.npy'),axis=0)
    theta__ = median(load('phi.npy'),axis=0)


    x0 = r_[f_max,pi,-2,0.8,theta_[:nch]]
   
    
    P, fopt, direc, _, _, warnflag = fmin_powell(Model, x0, args=(tvec,),maxiter=1e5,maxfun=1e6, xtol=1e-6, ftol=1e-6,full_output=True)
  
    fig = figure()
    title(fopt)

    Model(P,tvec,show_plt=True)
    fig.show()
    

    fig =figure('phase')
    plot(theta0[:16], P[4:20]-median(P[4:20]))
    plot(theta0[:16], theta_[:16],':')
    plot(theta0[:16], theta__[:16],'-.')

    if nch == 20:
        plot(theta0[16:], theta__[16:],'-.')
        plot(theta0[16:20], P[20:24],'--')

    xlim(0,2*pi)
    fig.show()
    pause(0.1)
    return P[4:]
 
    


def main():
    
    for path in ['graphs', 'mp3']:
        if not os.path.exists(path):
            os.mkdir(path)
            
    if len(sys.argv)==1 or sys.argv[1] ==  "plots":
        
        #tvec,sig = LoadSignal(12686)
        tvec,sig = LoadSignal2()

        hop = max((tvec[-1]- tvec[0])/800, 1e-5)
        multi_cwt(sig, 1e6, 5e-4, hop, tvec = tvec)
        
        from pygolem_lite.modules import saveconst

        os.system('convert -resize 150x120\! graphs/spectrogram0.png icon.png')
        saveconst('status', 0)





if __name__ == "__main__":
    main()