#!/usr/bin/python2
# -*- coding: utf-8 -*-
""" CREATED: 4/2013
AUTHOR: TOMÁŠ ODSTRČIL
"""
import matplotlib
matplotlib.rcParams['backend'] = 'Agg'
from scipy.signal import *
from numpy import *
from matplotlib.pyplot import *
from stft import stft
from matplotlib.ticker import MultipleLocator, FormatStrFormatter,LogLocator
import os
import sys
from pygolem_lite import load_adv, saveconst, Shot
from multiprocessing import Process, Pool, cpu_count
from scipy.stats.mstats import mquantiles
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 Spectrogram((frequenciCutOff,tvec, signal,name,start, end, plasma)):
print name
signal = double(signal)
startAdv = max(start*0.8, tvec[0]*1e3)
endAdv = min((end-start)*1.2+start, tvec[-1]*1e3)
# cut-off signal
ind = (tvec > startAdv*1e-3) & (tvec < endAdv*1e-3)
signal = signal[ind]
tvec = tvec[ind]
dt = mean(diff(tvec))
N = len(tvec)
for i in range(3):
filtered = medfilt(signal, 2*floor(N/500)+1)
res = abs(signal - filtered)
ind = argsort(res)
n_out = floor(N/400)
ind_2 = in1d(range(N), ind[-n_out:]) & (res > 0.2*amax(filtered))
signal[ind_2] = filtered[ind_2] #remove "n_out" most distant points from moving medianZ
print "computing spectrum"
#spec, scale = cwt(signal, dt, dj=0.05, wf='morlet', p=omega0, extmethod='none', extlength='powerof2')
fmin = 3000
fmax = frequenciCutOff
spec, freq = stft(signal, dt, dTime=1e-5, p=1000, extmethod='periodic', extlength='powerof2',res = 200,fmin=fmin,fmax=fmax)
print "done"
# approximate scales through frequencies
#freq = (omega0 + sqrt(2.0 + omega0**2))/(4*pi * scale[1:])
contrast = 10
#spec = spec[:,(tvec > startAdv )*( tvec <endAdv)]
pole = log(1+contrast*abs(spec))
#pole = pole[freq > frequenciCutOff, : ]
#freq = freq[freq > frequenciCutOff]
fig = figure()
ax = fig.add_axes([0.1, 0.1, 0.8, 0.85])
freq/=1000
img = ax.imshow(pole, extent=[startAdv,endAdv ,freq[-1], freq[0]], aspect='auto')
#ax.set_yscale('log', nonposy='clip')
minorLocator = MultipleLocator(1)
ax.xaxis.set_minor_locator(minorLocator)
ax.plot([start, start], [amin(freq), amax(freq)], 'w--')
ax.plot([end, end], [amin(freq), amax(freq)], 'w--')
ax.axis([startAdv,endAdv,amin(freq), amax(freq)])
ax.set_xlabel('time [ms]')
ax.set_ylabel('Frequency [kHz]')
savefig('data/spectrogram_'+name+'.png')
clf()
ind = (tvec > startAdv*1e-3 ) & ( tvec<endAdv*1e-3)
signal= signal[ ind ]
plot(1e3*tvec[ ind ],signal,'k',linewidth=0.1)
#show()
#axvline(linewidth=4)
#minimum = -max(abs(amin(signal)), abs(amax(signal)))
#maximum = -minimum
minimum = mquantiles(signal,0.01)
maximum = mquantiles(signal, 0.99)
minimum -= (maximum - minimum)*0.2
maximum += (maximum - minimum)*0.2
if maximum == minimum:
axis([startAdv,endAdv,None, None])
else:
axis([startAdv,endAdv,minimum, maximum])
#axis([startAdv,endAdv,None, None])
if plasma:
plot([start, start], [minimum, maximum], 'r--')
plot([end, end], [minimum, maximum], 'r--')
xlabel('time [ms]')
ylabel('Intensity [a.u.]')
#show()
savefig('data/signal_'+name+'.png')
clf()
Data = Shot()
plasma = Data['plasma']
if plasma:
start = Data['plasma_start']*1e3
end = Data['plasma_end']*1e3
else:
start = 0
end = 40
tvec, data = Shot()['nistandard6132']
tvec = float_(tvec)
frequenciCutOff = 8e4
try:
os.mkdir('data')
except:
pass
Ndim = size(data,1)
try:
p = Pool(cpu_count())
p.map( Spectrogram, [ (frequenciCutOff,tvec, data[:,i],str(i),start,end, plasma) for i in range(Ndim)])
p.close()
saveconst('status', 0)
except:
saveconst('status', 1)
raise
os.system('convert -resize 150 data/spectrogram_1.png icon.png')
