Source Code for Module spectractor.plotter

  1   
  2  '''This is Spectractor submodule containing classes for interactive actions 
  3  with matplotlib figures. 
  4   
  5  @author: Dmitry Nasonov 
  6  ''' 
  7   
  8   
  9  import matplotlib.pyplot as plt 
 10  from matplotlib.ticker import MultipleLocator, FormatStrFormatter 
 11  import numpy as np 
 12   
 13  from spectractor import Spliner 
 14   
 15   
 16 -class Clicker: 
 17      """Class for interactive manipulations with given subplot. 
 18      Coordinates of the cursor appends to list after click, plot updates with 
 19      cross and label. 
 20      @param splt: matplotlib subplot object 
 21      """ 
 22 -    def __init__(self, splt, xlims=None, clr='r', crdlim=30): 
 23          self.splt = splt 
 24          # Plot params 
 25          self.clr = clr 
 26          if xlims: 
 27              self.zw = abs(xlims[0]-xlims[1]) / 10. 
 28          else: 
 29              self.zw = None 
 30          self.zoomed = False 
 31          self.xlims = xlims 
 32          self.kpe, self.bpe = "key_press_event", "button_press_event" 
 33          self.tb = plt.get_current_fig_manager().toolbar 
 34          self.crdlim = crdlim 
 35          self.coords, self.plcm = [], [] 
 36   
 37 -    def getcoords(self, event, rndl=3): 
 38          return round(event.xdata, rndl), round(event.ydata, rndl) 
 39   
 40 -    def addpnt(self, event): 
 41          x, y = self.getcoords(event) 
 42          self.coords.append((x, y)) 
 43          if len(self.coords) > self.crdlim: 
 44              plt.close() 
 45          self.plcm.append(self.splt.plot(x, y, '+', c=self.clr, ms=10)[0]) 
 46          plt.draw() 
 47   
 48 -    def delplt(self, ind): 
 49          self.plcm[ind].remove() 
 50          del self.plcm[ind] 
 51          del self.coords[ind] 
 52          plt.draw() 
 53   
 54 -    def zoom(self, event): 
 55          if not self.zoomed: 
 56              x, y = self.getcoords(event) 
 57              x0, x1 = x-self.zw, x+self.zw 
 58              self.zoomed = True 
 59          else: 
 60              x0, x1 = self.xlims 
 61              self.zoomed = False 
 62          self.splt.set_xlim(x0, x1) 
 63          plt.draw() 
 64   
 65 -    def __call__(self, evnt): 
 66          """@param evnt: event""" 
 67          if not self.tb.mode and evnt.inaxes and evnt.name == self.bpe \ 
 68                  and evnt.button == 1: 
 69              self.addpnt(evnt) 
 70          elif evnt.name == self.kpe and evnt.key == ' ' and evnt.inaxes \ 
 71                  and self.zw: 
 72              self.zoom(evnt) 
 73          elif evnt.name == self.kpe and evnt.key == 'd' and len(self.coords): 
 74              self.delplt(-1) 
 75          elif evnt.name == self.kpe and evnt.key == 'q': 
 76              plt.close() 
 77   
 78   
 79 -class Flatter(Spliner): 
 80      """Class for interactive flatting set of 1D data. 
 81      @param splt: matplotlib subplot object 
 82      @param data: iterable array of 1D data 
 83      @param ccm: optional array of supplement reference points 
 84      @param blaze: optional array of supplement curves 
 85      @param zw: half width of zooming window 
 86      @param spk: degree of spline 
 87      """ 
 88 -    def __init__(self, splt, data, ccm=None, blaze=None, 
 89                  clr="c", lw=1, grid=False, zw=100, spk=3): 
 90          self.splt = splt 
 91          self.data = data 
 92          self.ccm = ccm 
 93          self.blz = blaze 
 94          # Plot params 
 95          self.clr = clr 
 96          self.lw = lw 
 97          self.grid = grid 
 98          self.zw = zw 
 99          self.kpe, self.bpe = "key_press_event", "button_press_event" 
100          self.tb = plt.get_current_fig_manager().toolbar 
101          self.nord, self.end = len(self.data), len(self.data[0]) 
102          self.xlims = (0, self.end) 
103          # Initialization 
104          self.Sl = Spliner(0, self.end, smooth=None, spk=spk) 
105          self.out = np.ones((self.nord, self.end)) 
106          self.init_prm() 
107          self.i = 0 
108          self.dat = self.data[0] 
109          self.plotdat() 
110   
111 -    def init_prm(self): 
112          self.plcrv, self.plspl, self.cmpl = [], [], [] 
113          self.coords = [] 
114          self.isdiv, self.zoomed = False, False 
115   
116 -    def plotdat(self): 
117          _tmp, = self.splt.plot(self.dat, '-k', lw=self.lw) 
118          self.plcrv.append(_tmp) 
119          if self.ccm is not None and self.i+1 in self.ccm: 
120              for x, y in self.ccm[self.i+1]: 
121                  self.cmpl.append(self.splt.plot(x, y, 'o', c='g', ms=5)[0]) 
122              self.splfit = self.mksplfit(self.ccm[self.i+1]) 
123              self.plspl.append(self.splt.plot(self.splfit, '-g')[0]) 
124          elif self.blz is not None: 
125              _tmp, = self.splt.plot(self.blz[self.i], '-g', lw=self.lw) 
126              self.plcrv.append(_tmp) 
127          self.set_lims(self.dat) 
128          self.splt.grid(self.grid) 
129          plt.draw() 
130   
131 -    def plotspl(self, delprev=True): 
132          self.splfit = self.mksplfit(self.coords) 
133          if delprev: 
134              [plelm.remove() for plelm in self.plspl] 
135              self.plspl = [] 
136          self.plspl.append(self.splt.plot(self.splfit, '-r')[0]) 
137          plt.draw() 
138   
139 -    def mksplfit(self, points): 
140          xfl, yfl = zip(*sorted(points)) 
141          return self.Sl.splpl(xfl, yfl) 
142   
143 -    def plotdiv(self): 
144          self.clear() 
145          self.dat = self.dat / self.splfit 
146          self.isdiv = True 
147          self.plcrv.append(self.splt.plot(self.dat, '-k', lw=self.lw)[0]) 
148          for x, y in self.coords: 
149              self.cmpl.append(self.splt.plot(x, 1, 'o', c=self.clr, ms=5)[0]) 
150          self.set_lims(self.dat) 
151          plt.draw() 
152   
153 -    def clear(self): 
154          [plelm.remove() for plelm in self.plcrv] 
155          [plelm.remove() for plelm in self.plspl] 
156          [plelm.remove() for plelm in self.cmpl] 
157          self.init_prm() 
158   
159 -    def delplt(self, ind): 
160          self.cmpl[ind].remove() 
161          del self.cmpl[ind] 
162          del self.coords[ind] 
163          plt.draw() 
164   
165 -    def set_lims(self, data, ysh=None): 
166          self.splt.set_xlim(self.xlims) 
167          dmin, dmax = np.min(data), np.max(data) 
168          if not ysh: 
169              ysh = (dmax - dmin) / 50. 
170          self.splt.set_ylim(dmin-ysh, dmax+ysh) 
171   
172 -    def getcoords(self, event, rndl=3): 
173          return round(event.xdata, rndl), round(event.ydata, rndl) 
174   
175 -    def addpnt(self, event): 
176          x, y = self.getcoords(event) 
177          if event.button == 3: 
178              y = self.dat[int(round(x))] 
179          self.coords.append((x, y)) 
180          self.cmpl.append(self.splt.plot(x, y, 'o', c=self.clr, ms=5)[0]) 
181          plt.draw() 
182   
183 -    def next(self): 
184          if self.isdiv: 
185              self.out[self.i] = self.dat 
186          self.clear() 
187          if self.i < self.nord - 1: 
188              self.i += 1 
189              self.dat = self.data[self.i] 
190              self.plotdat() 
191          else: 
192              plt.close() 
193   
194 -    def zoom(self, event): 
195          if not self.zoomed: 
196              x, y = self.getcoords(event) 
197              x0, x1 = max(x-self.zw, 0), min(self.end, x+self.zw) 
198              self.zoomed = True 
199          else: 
200              x0, x1 = 0, self.end 
201              self.zoomed = False 
202          self.xlims = (x0, x1) 
203          self.set_lims(self.dat[x0:x1]) 
204          plt.draw() 
205   
206 -    def __call__(self, evnt): 
207          """@param evnt: event""" 
208          if not self.tb.mode and evnt.inaxes and evnt.name == self.bpe \ 
209                  and evnt.button in (1, 3): 
210              self.addpnt(evnt) 
211          elif evnt.name == self.kpe and evnt.inaxes and evnt.key == ' ': 
212              self.zoom(evnt) 
213          elif evnt.name == self.kpe and evnt.key == 'f10' \ 
214                  and len(self.coords) > 3: 
215              self.plotspl() 
216          elif evnt.name == self.kpe and evnt.key == 'f9': 
217              self.plotdiv() 
218          elif evnt.name == self.kpe and evnt.key == 'd' and len(self.coords): 
219              self.delplt(-1) 
220          elif evnt.name == self.kpe and evnt.key == 'D' and len(self.coords): 
221              self.delplt(0) 
222          elif evnt.name == self.kpe and evnt.key == 'n': 
223              self.next() 
224          elif evnt.name == self.kpe and evnt.key == 'q': 
225              plt.close() 
226   
227   
228 -def plt_opt(figsub, **kwargs): 
229      """Set default parameters for kosher plotting""" 
230      lblsz = 'x-small' 
231      if 'lblsz' in kwargs: 
232          lblsz = kwargs['lblsz'] 
233      if 'xlbl' in kwargs: 
234          figsub.set_xlabel(kwargs['xlbl'], fontsize=lblsz) 
235      if 'ylbl' in kwargs: 
236          figsub.set_ylabel(kwargs['ylbl'], fontsize=lblsz) 
237      if 'grid' in kwargs: 
238          figsub.grid(True) 
239      if 'xlim' in kwargs: 
240          figsub.set_xlim(kwargs['xlim']) 
241      if 'ylim' in kwargs: 
242          figsub.set_ylim(kwargs['ylim']) 
243      if 'xlc' in kwargs: 
244          figsub.xaxis.set_major_locator(MultipleLocator(kwargs['xlc'][0])) 
245          figsub.xaxis.set_minor_locator(MultipleLocator(kwargs['xlc'][1])) 
246      if 'ylc' in kwargs: 
247          figsub.yaxis.set_major_locator(MultipleLocator(kwargs['ylc'][0])) 
248          figsub.yaxis.set_minor_locator(MultipleLocator(kwargs['ylc'][1])) 
249      if 'xmjf' in kwargs: 
250          figsub.xaxis.set_major_formatter(FormatStrFormatter(kwargs['xmjf'])) 
251      if 'ticksz' in kwargs: 
252          #figsub.set_xticklabels(fontsize=kwargs['ticksz']) 
253          figsub.ticklabel_format(fontsize=kwargs['ticksz']) 
254      if 'legend' in kwargs: 
255          leglblsz = 'x-small' 
256          if 'leglblsz' in kwargs: 
257              leglblsz = kwargs['leglblsz'] 
258          leg = figsub.legend(numpoints=1, loc=kwargs['loc'], #'best' 
259                      labelspacing=kwargs['legend']) 
260          ltext = leg.get_texts() 
261          plt.setp(ltext, fontsize=leglblsz) 
262