| 1 | from asap import scantable
|
|---|
| 2 | from pylab import *
|
|---|
| 3 | from numpy import array, ma, logical_not
|
|---|
| 4 | from numpy.fft import fft, ifft
|
|---|
| 5 |
|
|---|
| 6 | class lagplotter(object):
|
|---|
| 7 | def __init__(self, stable):
|
|---|
| 8 | if not isinstance(stable, scantable):
|
|---|
| 9 | raise TypeError("This works only on scantables")
|
|---|
| 10 | self._scan = stable
|
|---|
| 11 | self.figure = None
|
|---|
| 12 | self.dataaxes = None
|
|---|
| 13 | self.fftaxes = None
|
|---|
| 14 | self.resultaxes = None
|
|---|
| 15 | self.flags = []
|
|---|
| 16 |
|
|---|
| 17 | def _init_plot(self):
|
|---|
| 18 | if not self.figure:
|
|---|
| 19 | self.figure = figure()
|
|---|
| 20 | self.dataaxes = self.figure.add_subplot(3,1,1)
|
|---|
| 21 | self.fftaxes = self.figure.add_subplot(3,1,2)
|
|---|
| 22 | self.resultaxes = self.figure.add_subplot(3,1,3)
|
|---|
| 23 | else:
|
|---|
| 24 | self.dataaxes.cla()
|
|---|
| 25 | self.fftaxes.cla()
|
|---|
| 26 | self.resultaxes.cla()
|
|---|
| 27 |
|
|---|
| 28 |
|
|---|
| 29 | def flag(self):
|
|---|
| 30 | self._init_plot()
|
|---|
| 31 | for i in xrange(len(self._scan)):
|
|---|
| 32 | self.dataaxes.cla()
|
|---|
| 33 | self.fftaxes.cla()
|
|---|
| 34 | self.resultaxes.cla()
|
|---|
| 35 | x = array(self._scan.get_abcissa(i)[0])
|
|---|
| 36 | y = array(self._scan.get_spectrum(i))
|
|---|
| 37 | msk = self._scan.get_mask(i)
|
|---|
| 38 | marr = ma.MaskedArray(y, logical_not(msk), fill_value=0.0)
|
|---|
| 39 | self.dataaxes.plot(x, marr)
|
|---|
| 40 | nfft = len(marr)/2+1
|
|---|
| 41 | yfft = fft(marr.filled())
|
|---|
| 42 | self.fftaxes.semilogy(abs(yfft)[0:nfft])
|
|---|
| 43 | self.figure.show()
|
|---|
| 44 | raw_input("Press any key to continue...")
|
|---|
| 45 | flagstart = int(ginput(show_clicks=False)[0][0]+0.5)
|
|---|
| 46 | flagend = int(ginput(show_clicks=False)[0][0]+0.5)
|
|---|
| 47 | xfft = range(len(yfft))
|
|---|
| 48 | self.fftaxes.axvspan(flagstart, flagend, alpha=0.3)
|
|---|
| 49 | self.fftaxes.axvspan(xfft[-flagend], xfft[-flagstart], alpha=0.3)
|
|---|
| 50 | yfft[flagstart:flagend] = 0+0j
|
|---|
| 51 | yfft[-flagend:-flagstart] = 0+0j
|
|---|
| 52 | yi = ifft(yfft)
|
|---|
| 53 | self.resultaxes.plot(x, yi)
|
|---|
| 54 | self.figure.show()
|
|---|
| 55 | approve=raw_input("Commit flags y/n ? ")
|
|---|
| 56 | if approve in ["y", "Y"]:
|
|---|
| 57 | self.flags.append([flagstart, flagend])
|
|---|
| 58 | self._scan.set_spectrum(yi.real, i)
|
|---|
| 59 | else:
|
|---|
| 60 | del self.fftaxes.patches[-1]
|
|---|
| 61 | cont = raw_input("Continue (c) or quit (q)? ")
|
|---|
| 62 | if not cont == "c":
|
|---|
| 63 | return self.flags
|
|---|
| 64 | return self.flags
|
|---|
