Source Code for Module reflectometry.reduction.smoothcor

  1  """ 
  2  Data smoothing 
  3   
  4  Uses moving window 1-D polynomial least squares smoothing filter. 
  5   
  6  Usage 
  7  ===== 
  8   
  9  Create and apply the filter:: 
 10   
 11      from reflectometry.reduction.corrections import smooth 
 12      ... 
 13      data.apply(smooth(degree=2,span=12)) 
 14   
 15  """ 
 16  __all__ = ['Smooth'] 
 17   
 18  import numpy 
 19   
 20  from reflectometry.reduction.correction import Correction 
 21  from reflectometry.reduction.wsolve import wpolyfit 
 22   
 23   
 24 -class Smooth(Correction): 
 25      """ 
 26      Moving window 1-D polynomial least squares smoothing filter. 
 27   
 28      The parameters are the polynomial order and the window size. 
 29      The window size is the number of consecutive points used to 
 30      smooth the data.  The window size must be odd. 
 31      """ 
 32   
 33      properties = ['degree','span'] 
 34      degree = 2 
 35      """Polynomial order for smoothing""" 
 36      span = 11 
 37      """Number of points used to fit smoothing polynomial""" 
 38   
 39 -    def __init__(self, **kw): 
 40          """ 
 41          Define the smoothing polynomial. 
 42          """ 
 43          for k,v in kw.iteritems(): 
 44              assert hasattr(self,k), "No %s in %s"%(k,self.__class__.__name__) 
 45              setattr(self,k,v) 
 46          assert self.span%2==1, "Span must be odd" 
 47   
 48 -    def apply(self, data): 
 49          """Apply the correction to the data""" 
 50          if data.ispolarized(): 
 51              lines = [data.pp, data.pm, data.mp, data.mm] 
 52          else: 
 53              lines = [data] 
 54          for L in lines: 
 55              v,dv = smooth(L.x, L.v, L.dv, degree=self.degree, span=self.span) 
 56              L.v,L.dv = v,dv 
 57   
 58 -    def __str__(self): 
 59          return "Smooth(degree=%d,span=%d)"%(self.degree,self.span) 
 60   
 61  # TODO: check for equal spacing and use Savitsky-Golay since it will 
 62  # TODO: be orders of magnitude faster. Precompute the SG coefficients 
 63  # TODO: in the Smooth class. 
 64   
 65   
 66 -def smooth(x, y, dy=None, degree=2, span=5): 
 67      """ 
 68      Moving least squares smoother. 
 69   
 70      If x is equally spaced, then this is equivalent to a Savitsky-Golay 
 71      filter of the same degree and span. 
 72      """ 
 73      if dy == None: dy = numpy.ones(y.shape) 
 74      yp,dyp = numpy.empty(y.shape),numpy.empty(y.shape) 
 75   
 76      n = len(x) 
 77      k = (span-1)/2 
 78   
 79      poly = wpolyfit(x[:span], y[:span], dy[:span], degree=degree) 
 80      yp[:k+1],dyp[:k+1] = poly.ci(x[:k+1]) 
 81   
 82      for i in range(k+1,n-k+1): 
 83          poly = wpolyfit(x[i-k:i+k],y[i-k:i+k],dy[i-k:i+k], degree=degree) 
 84          yp[i:i+1],dyp[i:i+1] = poly.ci(x[i:i+1]) 
 85   
 86      poly = wpolyfit(x[-span:],y[-span:],dy[-span:],degree=degree) 
 87      yp[-k-1:],dyp[-k-1:] = poly.ci(x[-k-1:]) 
 88   
 89      return yp,dyp 
 90   
 91   
 92 -def demo(): 
 93      import pylab 
 94      from reflectometry.reduction.examples import e3a12 as dataset 
 95      data = dataset.slits() 
 96      pylab.errorbar(data.pp.x,data.pp.v,data.pp.dv,fmt='xg') 
 97      data.apply(Smooth(degree=2,span=7)) 
 98      pylab.semilogy(data.pp.x, data.pp.v, '-g', 
 99                     data.pp.x, data.pp.v-data.pp.dv, '-.g', 
100                     data.pp.x, data.pp.v+data.pp.dv, '-.g') 
101      pylab.show() 
102   
103  if __name__ == "__main__": demo() 
104