Source code :: main

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

import matplotlib
matplotlib.rcParams['backend'] = 'Agg'
matplotlib.rc('font',  size='10')
matplotlib.rc('text', usetex=True)  # FIXME !! nicer but slower !!!
from matplotlib.pyplot import *
import os

from scipy.signal import fftconvolve
from multiprocessing import Process, Pool, cpu_count
import time
from numpy import *
from CMWT import *
from SoundGenerator import *
from pygolem_lite import Shot
from pygolem_lite.modules import list2array,deconvolveExp,save_adv,saveconst
from scipy.stats.mstats  import mquantiles


RingCoilOrientation = (1,1,-1,-1, 1,-1,1,1,1,-1,-1,1,1,1,1,-1) #TODO důležité
MirnovCoil = (1,1,1,1)#TODO důležité


matplotlib.rcParams['xtick.direction'] = 'out'
matplotlib.rcParams['ytick.direction'] = 'out'

matplotlib.rcParams['xtick.major.size'] = 10
matplotlib.rcParams['xtick.minor.size'] = 7

matplotlib.rcParams['ytick.major.size'] = 10
matplotlib.rcParams['ytick.minor.size'] = 7



def PrepareData(data,typ):
    sign = MirnovCoil   #WARNING závisí na aktuálním nastavení cívek!!TODO dát to do externího configu
    n = size(data, 0)
    envelope = zeros(shape(data))
    
    dt = data[0,1]-data[0,0]
    n_smooth = int(0.002/dt)*2+1  #omezení časového rozlišení pro nízké frekvence 250Hz
    
    for i in range(1,n): 
	data[i,:]*= sign[i-1]
	baseline  = fftconvolve((data[i,:]), ones(n_smooth)/n_smooth,mode = 'same' )
	data[i,:]-= baseline
	envelope[i,:] = fftconvolve(abs(data[i,:]), ones(n_smooth/10)/n_smooth/10,mode = 'same' )
	plot(baseline, label = str(i))
    
    legend()
    #show()
    savefig('data.png')
    close()
    total_envelope  = mean(envelope,0)
    for i in range(1,n):    
	data[i,:]*= total_envelope/(envelope[i,:]+1e-6)

    return data
  
      
def plotData(  data):
    dt = data[0,1]-data[0,0]

    n_smooth = 20  #omezení časového rozlišení pro nízké frekvence 250Hz
    n = size(data, 1)
    for i in range(1,size(data, 0)):
	#baseline  = fftconvolve((data[i,:]), ones(n_smooth)/n_smooth,mode = 'same' )
	dfft = fft.rfft(data[i,:])
	dfft[n/16:] = 0
	difft = fft.irfft(dfft)
	#plot(data[0,:], data[i,:])

	plot(data[0,:], difft)
	xlim(0.016, 0.0171)
	ylim(-0.2, 0.2)
    #show()
	
    savefig('./graphs/raw_data.png')
    close()
	
def loadconst(fname):
    with open(fname, 'r') as fhandle:
	return float(fhandle.readline()) #return the raw string   	

	
def LoadData():
    Data = Shot()

    gd = Shot().get_data
    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)

    Papouch = list2array( Data[das, [m1, m5, m9, m13]] ).T 
      
    plasma = Data['plasma']

    plasma_start = Data['plasma_start']
    plasma_end = Data['plasma_end']

    return Papouch,plasma_start, plasma_end,plasma
    
    
#def LoadData(path, shot_num, skip_det_num):
    
    #path += str(shot_num)+'/'
    #try:
	#data = load(path+'data.npy')
	#if size(data,1) == 0:
	    #raise "wrong data"
    #except:
	##print 'reload'
	#data = list()
	#tvec = None
	#for i in range(0,16):
	    #print 'load NIturbo_%2.2i' %(i+1)

	    #if i in skip_det_num:
		#data.append(zeros(shape(data[0])))
		#continue
	    ##single_data =  loadtxt(path+'NIturbo_%2.2i' %(i+1)+'.asc', usecols = (1,))
	    #single_data =  loadtxt(path+'NIturbo_%2.2i' %(i+1), usecols = (1,))
	    
	    #if tvec == None:
		##tvec =  loadtxt(path+'NIturbo_%2.2i' %(i+1)+'.asc', usecols = (0,))	
		#tvec =  loadtxt(path+'NIturbo_%2.2i' %(i+1), usecols = (0,))	

	    #data.append(single_data)
	    
	    
	#data = array(data)
	#if size(data,1) == 0:
	    #raise "wrong data"
	#data = vstack((tvec, data))
	#save(path+'data', single(data))
    
    #try:
	#start = loadconst(path+'PlasmaStart')/1000
	#end   = loadconst(path+'PlasmaEnd'  )/1000
    #except:
	#start = nan
	#end = nan
    #print start,end
    #return data,start,end
	
    
#def LoadData_npz(path, shot_num, skip_det_num):  #FIXME provizorní verze!!
    #data = load('Nidatap.npz')
    ##data  data['data']
    
    #tvec = linspace(data['t_start'], data['t_end'], size(data['data'],0))
    #data = data['data']
    #data[:,skip_det_num] = 0
    #data = vstack((tvec, data.T))

    #start = data[0,8300]
    #end = data[0,19260]
    #return data,start,end

##def LoadData_mirnov(path, shot_num):
    
    ##path += str(shot_num)+'/'
    ##try:
	##data = load(path+'data.npy')
	##if size(data,1) == 0:
	    ##raise "wrong data"
	##if shot_num > 6000 and shot_num < 9500:
	    ##data = data.T
	    ##data[1:,1:] = diff(data[1:,:], axis = 1)/(data[0,1]-data[0,0])
    ##except:
	##data = list()
	##tvec = None
	##for i in range(0,4):
	    ##print 'load PapouchSt_%2.2i' %(i+1)


	    ##single_data =  loadtxt(path+'PapouchSt_%2.2i' %(i+1), usecols = (1,))
	    
	    ##if tvec == None:
		##tvec =  loadtxt(path+'PapouchSt_%2.2i' %(i+1), usecols = (0,))	

	    ##data.append(single_data)
	    
	    
	##data = array(data)
	##if size(data,1) == 0:
	    ##raise "wrong data"
	##data = vstack((tvec, data))
	##save(path+'data', single(data))
    

    ##try:
	##start = loadconst(path+'PlasmaStart')/1000
	##end   = loadconst(path+'PlasmaEnd'  )/1000
    ##except:
	##start = nan
	##end = nan
    ##print start,end
    ##return data,start,end    
    
def PlotSpec((freq, field, t,  sufname ,vmin, vmax)):
    
    fig = figure()
    t = t*1000
    ax = fig.add_axes([0.1, 0.1, 0.8, 0.85])
    ax.set_yscale('log', nonposy='clip')

    img = ax.imshow(abs(field), extent=[t[0],t[-1] ,freq[-1], freq[0]], aspect='auto',vmin=vmin, vmax=vmax) 
    minorLocator   = MultipleLocator(1)
    ax.xaxis.set_minor_locator(minorLocator)
    ax.axis([t[0],t[-1] ,amin(freq), amax(freq)])

    ax.set_xlabel('time [ms]')


    ax.set_ylabel('Frequency [Hz]')
    savefig('graphs/spectrogram'+str(sufname)+'.png',bbox_inches='tight')
    close()
    
def CalculateSpectrogram():
      
    print "CalculateSpectrogram"
    
    t1 = time.time()

    signal,plasma_start,plasma_end,plasma = LoadData()
    
    #plasma = False  #BUG!!!!
    
    
    if not plasma:
	plasma_start = 0
	plasma_end = 0.04
    
    dt = signal[0,1]-signal[0,0]
    n_start = int((plasma_start - signal[0,0])/dt)
    n_end= int((plasma_end -signal[0,0] )/dt)

    
    t = signal[0,n_start:n_end]
    signal = signal[1:,n_start:n_end].T

    omega0 = 10#20
    horiz_res = 2000
    f_min = 1e3  #Hz
    f_max = 100e3 #Hz
    
    N = size(signal,0)
    n_det = size(signal,1)
    t2= time.time()
    signal_fft = fft.rfft(signal, axis = 0)
    for m in arange(-2,3):
	print "wave order", m
	signal_m = copy(signal_fft)
	phase = arange(n_det)/double(n_det)*m 
	exp_phase = matrix(exp(-2*pi*1j*phase))
	signal_m*= exp_phase
	signal_m = sum(signal_m, axis = 1)

    
	signal_m = fft.irfft(signal_m)
	try:
	  soundGenerator(signal_m,2000, 'mp3/sound'+str(m))
	except Exception, e:
	  print "sound gener. failed err:" , e.message
	  
    print 'generated sound', time.time()-t2


    # !! searched MHD mdoes !! 
    modes =  [-1,0,1,2]
    spec_all = list()
    scale_all = list()
	
    print "started wavelets"
    
    p = Pool(cpu_count())
    out = p.map(NTM_CWT, [ (signal, dt, 0.005,omega0,m ,horiz_res,  f_min,f_max) for m in modes ])
    p.close()
    p.join()

    for i in range(len(modes)):
	spec, scale = out[i]
	spec = single(abs(spec))
	spec_all.append(spec)
	scale_all.append(scale)
	
    print 'calc time',  time.time()-t1


    t_plot = time.time()
    spec_all  = array(spec_all)
    scale_all = array(scale_all)

    freq = (omega0 + sqrt(2.0 + omega0**2))/(4*pi * scale_all)
    contrast = 10
    spec_all = log(1+contrast*abs(spec_all))
    vmin = amin(spec_all)
    vmax = amax(spec_all)
    p = Pool(cpu_count())

    p.map( PlotSpec, [(freq[i,...],spec_all[i,...],t, m,vmin, vmax) for i,m in enumerate(modes) ])
    p.close()
    p.join()

    print 'plot. time', time.time()-t_plot




def main():
    
    for path in ['graphs', 'mp3']:
	if not os.path.exists(path):
	    os.mkdir(path)
	    

    if sys.argv[1] ==  "plots":
	CalculateSpectrogram()
	os.system('convert -resize 150x120\! graphs/spectrogram0.png icon.png')
	saveconst('status', 0)





if __name__ == "__main__":
    main()