from asap import scantable from pylab import * from numpy import array, ma, logical_not from numpy.fft import fft, ifft class lagplotter(object): def __init__(self, stable): if not isinstance(stable, scantable): raise TypeError("This works only on scantables") self._scan = stable self.figure = None self.dataaxes = None self.fftaxes = None self.resultaxes = None self.flags = [] def _init_plot(self): if not self.figure: self.figure = figure() self.dataaxes = self.figure.add_subplot(3,1,1) self.fftaxes = self.figure.add_subplot(3,1,2) self.resultaxes = self.figure.add_subplot(3,1,3) else: self.dataaxes.cla() self.fftaxes.cla() self.resultaxes.cla() def flag(self): self._init_plot() for i in xrange(len(self._scan)): self.dataaxes.cla() self.fftaxes.cla() self.resultaxes.cla() x = array(self._scan.get_abcissa(i)[0]) y = array(self._scan.get_spectrum(i)) msk = self._scan.get_mask(i) marr = ma.MaskedArray(y, logical_not(msk), fill_value=0.0) self.dataaxes.plot(x, marr) nfft = len(marr)/2+1 yfft = fft(marr.filled()) self.fftaxes.semilogy(abs(yfft)[0:nfft]) self.figure.show() raw_input("Press any key to continue...") print "Now select a start and end point by clicking on the middle plot" print "Start point ..." flagstart = int(ginput(show_clicks=False)[0][0]+0.5) print "End point ..." flagend = int(ginput(show_clicks=False)[0][0]+0.5) xfft = range(len(yfft)) self.fftaxes.axvspan(flagstart, flagend, alpha=0.3) yfft[flagstart:flagend] = 0+0j yfft[-flagend:-flagstart] = 0+0j yi = ifft(yfft) self.resultaxes.plot(x, yi) self.figure.show() inp = raw_input("Commit flags (c), keep (k) or ignore(i)? ")\ .lower() if inp.startswith("c") self.flags.append([flagstart, flagend]) self._scan.set_spectrum(yi.real, i) elif inp.startswith("k"): self.flags.append([flagstart, flagend]) else: del self.fftaxes.patches[-1] cont = raw_input("Continue (c) or quit (q)? ") if not cont == "c": return self.flags return self.flags