Source Code for Module linear_operators.wrappers.pywt_lo

 1  """ Wrap pywt transform into LinearOperators """ 
 2  import numpy as np 
 3  import pywt 
 4  from ..interface import LinearOperator 
 5  from ..ndoperators import NDOperator 
 6  from copy import copy 
 7   
 8 -def wavedec(shapein, wavelet, mode='sym', level=None, dtype=np.float64): 
 9      a = np.zeros(shapein) 
10      b = pywt.wavedec(a, wavelet, mode=mode, level=level) 
11      def matvec(x): 
12          dec_list = pywt.wavedec(x, wavelet, mode=mode, level=level) 
13          return np.concatenate(dec_list) 
14      def rmatvec(x): 
15          x_list = [] 
16          count = 0 
17          for el in b: 
18              n_el = np.asarray(el).size 
19              x_list.append(np.array(x[count:count+n_el])) 
20              count += n_el 
21          return pywt.waverec(x_list, wavelet, mode=mode) 
22      shapeout = matvec(a).size 
23      return lo.LinearOperator((shapeout, shapein), matvec=matvec, rmatvec=rmatvec, 
24                               dtype=dtype) 
25   
26 -def wavelet2(shapein, wavelet, mode='zpd', level=None, dtype=np.float64): 
27      """ 
28      2d wavelet decomposition / reconstruction as a NDOperator 
29      """ 
30      a = np.zeros(shapein) 
31      b = pywt.wavedec2(a, wavelet, mode=mode, level=level) 
32      shapeout = coefs2array(b).shape 
33      def matvec(x): 
34          coefs = pywt.wavedec2(x, wavelet, mode=mode, level=level) 
35          return coefs2array(coefs) 
36      def rmatvec(x): 
37          coefs = array2coefs(x, b) 
38          return pywt.waverec2(coefs, wavelet, mode=mode)[:shapein[0], :shapein[1]] 
39      return lo.NDOperator(shapein, shapeout, matvec, rmatvec, dtype=dtype) 
40   
41 -def coefs2array(coefs): 
42      out = coefs[0] 
43      for scale in coefs[1:]: 
44          if out.shape[0] == scale[0].shape[0] + 1: 
45              out = out[:-1] 
46          if out.shape[1] == scale[0].shape[1] + 1: 
47              out = out[:, :-1] 
48          out = np.vstack((np.hstack((out, scale[0])), np.hstack((scale[1], scale[2])))) 
49      return out 
50   
51 -def array2coefs(a, l): 
52      ilim = [0, l[0].shape[0]] 
53      jlim = [0, l[0].shape[1]] 
54      coefs = [a[ilim[0]:ilim[1], jlim[0]:jlim[1]], ] 
55      for i in xrange(1, len(l)): 
56          ilim.append(ilim[-1] + l[i][0].shape[0]) 
57          jlim.append(jlim[-1] + l[i][0].shape[1]) 
58          scale = (a[0:ilim[-2], jlim[-2]:jlim[-1]], 
59                   a[ilim[-2]:ilim[-1], 0:jlim[-2]], 
60                   a[ilim[-2]:ilim[-1], jlim[-2]:jlim[-1]]) 
61          coefs.append(scale) 
62      return coefs 
63