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 |
---|