#!/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 = True ring=False 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'] print tvec out = PrepareSignal(tvec, data,tstart,tend,shot) return out 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)) if nstart == nend: return tvec, data 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()