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