Python pywt.coeffs_to_array() Examples

def test_coeffs_to_array():
    # single element list returns the first element
    a_coeffs = [np.arange(8).reshape(2, 4), ]
    arr, arr_slices = pywt.coeffs_to_array(a_coeffs)
    assert_allclose(arr, a_coeffs[0])
    assert_allclose(arr, arr[arr_slices[0]])

    assert_raises(ValueError, pywt.coeffs_to_array, [])
    # invalid second element:  array as in wavedec, but not 1D
    assert_raises(ValueError, pywt.coeffs_to_array, [a_coeffs[0], ] * 2)
    # invalid second element:  tuple as in wavedec2, but not a 3-tuple
    assert_raises(ValueError, pywt.coeffs_to_array, [a_coeffs[0],
                                                     (a_coeffs[0], )])
    # coefficients as None is not supported
    assert_raises(ValueError, pywt.coeffs_to_array, [None, ])
    assert_raises(ValueError, pywt.coeffs_to_array, [a_coeffs,
                                                     (None, None, None)])

    # invalid type for second coefficient list element
    assert_raises(ValueError, pywt.coeffs_to_array, [a_coeffs, None])

    # use an invalid key name in the coef dictionary
    coeffs = [np.array([0]), dict(d=np.array([0]), c=np.array([0]))]
    assert_raises(ValueError, pywt.coeffs_to_array, coeffs) 

def test_wavedecn_coeff_reshape_even():
    # verify round trip is correct:
    #   wavedecn - >coeffs_to_array-> array_to_coeffs -> waverecn
    # This is done for wavedec{1, 2, n}
    rng = np.random.RandomState(1234)
    params = {'wavedec': {'d': 1, 'dec': pywt.wavedec, 'rec': pywt.waverec},
              'wavedec2': {'d': 2, 'dec': pywt.wavedec2, 'rec': pywt.waverec2},
              'wavedecn': {'d': 3, 'dec': pywt.wavedecn, 'rec': pywt.waverecn}}
    N = 28
    for f in params:
        x1 = rng.randn(*([N] * params[f]['d']))
        for mode in pywt.Modes.modes:
            for wave in wavelist:
                w = pywt.Wavelet(wave)
                maxlevel = pywt.dwt_max_level(np.min(x1.shape), w.dec_len)
                if maxlevel == 0:
                    continue

                coeffs = params[f]['dec'](x1, w, mode=mode)
                coeff_arr, coeff_slices = pywt.coeffs_to_array(coeffs)
                coeffs2 = pywt.array_to_coeffs(coeff_arr, coeff_slices,
                                               output_format=f)
                x1r = params[f]['rec'](coeffs2, w, mode=mode)

                assert_allclose(x1, x1r, rtol=1e-4, atol=1e-4) 

def test_coeffs_to_array_padding():
    rng = np.random.RandomState(1234)
    x1 = rng.randn(32, 32)
    mode = 'symmetric'
    coeffs = pywt.wavedecn(x1, 'db2', mode=mode)

    # padding=None raises a ValueError when tight packing is not possible
    assert_raises(ValueError, pywt.coeffs_to_array, coeffs, padding=None)

    # set padded values to nan
    coeff_arr, coeff_slices = pywt.coeffs_to_array(coeffs, padding=np.nan)
    npad = np.sum(np.isnan(coeff_arr))
    assert_(npad > 0)

    # pad with zeros
    coeff_arr, coeff_slices = pywt.coeffs_to_array(coeffs, padding=0)
    assert_(np.sum(np.isnan(coeff_arr)) == 0)
    assert_(np.sum(coeff_arr == 0) == npad)

    # Haar case with N as a power of 2 can be tightly packed
    coeffs_haar = pywt.wavedecn(x1, 'haar', mode=mode)
    coeff_arr, coeff_slices = pywt.coeffs_to_array(coeffs_haar, padding=None)
    # shape of coeff_arr will match in this case, but not in general
    assert_equal(coeff_arr.shape, x1.shape) 

def test_waverecn_coeff_reshape_odd():
    # verify round trip is correct:
    #   wavedecn - >coeffs_to_array-> array_to_coeffs -> waverecn
    rng = np.random.RandomState(1234)
    x1 = rng.randn(35, 33)
    for mode in pywt.Modes.modes:
        for wave in ['haar', ]:
            w = pywt.Wavelet(wave)
            maxlevel = pywt.dwt_max_level(np.min(x1.shape), w.dec_len)
            if maxlevel == 0:
                continue
            coeffs = pywt.wavedecn(x1, w, mode=mode)
            coeff_arr, coeff_slices = pywt.coeffs_to_array(coeffs)
            coeffs2 = pywt.array_to_coeffs(coeff_arr, coeff_slices)
            x1r = pywt.waverecn(coeffs2, w, mode=mode)
            # truncate reconstructed values to original shape
            x1r = x1r[[slice(s) for s in x1.shape]]
            assert_allclose(x1, x1r, rtol=1e-4, atol=1e-4) 

def hfilter(diff_image, var_image, threshold=1, ndamp=10):
    """
    This code was inspired from: https://github.com/spacetelescope/sprint_notebooks/blob/master/lucy_damped_haar.ipynb
    I believe it was initially written by Justin Ely: https://github.com/justincely
    It was buggy and not working properly with every image sizes.
    I have thus exchanged it by using pyWavelet (pywt) and a custom function htrans
    to calculate the matrix for the var_image.
    """
    him, coeff_slices = pywt.coeffs_to_array(pywt.wavedec2(diff_image.astype(np.float), 'haar'), padding=0)
    dvarim = htrans(var_image.astype(np.float))
    
    sqhim = ((him/threshold)**2)/dvarim
    index = np.where(sqhim < 1)
    
    if len(index[0]) == 0:
        return diff_image
    
    # Eq. 8 of White is derived leading to N*x^(N-1)-(N-1)*x^N  :DOI: 10.1117/12.176819
    sqhim = sqhim[index] * (ndamp * sqhim[index]**(ndamp-1) - (ndamp-1)*sqhim[index]**ndamp)
    him[index] = sign(threshold*np.sqrt(dvarim[index] * sqhim), him[index])
    
    return pywt.waverec2(pywt.array_to_coeffs(him, coeff_slices, output_format='wavedec2'), 'haar')[:diff_image.shape[0],:diff_image.shape[1]] 

def __init__(self,nrow=256,ncol=256,wavelet='db4',level=3,fwd_mode='recon',\
        dtype=np.float64,name=None):

        # Save parameters
        self.wavelet = wavelet
        self.level = level
        shape0 = (nrow,ncol)
        shape1 = (nrow,ncol)
        dtype0 = dtype
        dtype1 = dtype

        if pywt.Wavelet(wavelet).orthogonal:
            svd_avail = True #SVD calculation assumes an orthogonal wavelet
        else:
            svd_avail = False
        BaseLinTrans.__init__(self, shape0, shape1, dtype0, dtype1,\
           svd_avail=svd_avail,name=name)


        # Set the mode to periodic to make the wavelet orthogonal
        self.mode = 'periodization'

        # Send a zero image to get the coefficient slices
        im = np.zeros((nrow,ncol))
        coeffs = pywt.wavedec2(im, wavelet=self.wavelet, level=self.level, \
            mode=self.mode)
        _, self.coeff_slices = pywt.coeffs_to_array(coeffs)


        # Confirm that fwd_mode is valid
        if (fwd_mode != 'recon') and (fwd_mode != 'analysis'):
            raise common.VpException('fwd_mode must be recon or analysis')
        self.fwd_mode = fwd_mode 

def analysis(self,z0):
        """
        Analysis:  image -> coefficients
        """
        coeffs = pywt.wavedec2(z0, wavelet=self.wavelet, level=self.level, \
            mode=self.mode)
        z1, _ = pywt.coeffs_to_array(coeffs)
        return z1 

def __init__(self, dims, dirs=(0, 1), wavelet='haar',
                 level=1, dtype='float64'):
        if pywt is None:
            raise ModuleNotFoundError('The wavelet operator requires '
                                      'the pywt package t be installed. '
                                      'Run "pip install PyWavelets" or '
                                      '"conda install pywavelets".')
        _checkwavelet(wavelet)

        # define padding for length to be power of 2
        ndimpow2 = [max(2 ** ceil(log(dims[dir], 2)), 2 ** level)
                    for dir in dirs]
        pad = [(0, 0)] * len(dims)
        for i, dir in enumerate(dirs):
            pad[dir] = (0, ndimpow2[i] - dims[dir])
        self.pad = Pad(dims, pad)
        self.dims = dims
        self.dirs = dirs
        self.dimsd = list(dims)
        for i, dir in enumerate(dirs):
            self.dimsd[dir] = ndimpow2[i]

        # apply transform once again to find out slices
        _, self.sl = \
            pywt.coeffs_to_array(pywt.wavedec2(np.ones(self.dimsd),
                                               wavelet=wavelet,
                                               level=level,
                                               mode='periodization',
                                               axes=self.dirs),
                                 axes=self.dirs)
        self.wavelet = wavelet
        self.waveletadj = _adjointwavelet(wavelet)
        self.level = level
        self.shape = (int(np.prod(self.dimsd)), int(np.prod(self.dims)))
        self.dtype = np.dtype(dtype)
        self.explicit = False 

def _matvec(self, x):
        x = self.pad.matvec(x)
        x = np.reshape(x, self.dimsd)
        y = pywt.coeffs_to_array(pywt.wavedec2(x, wavelet=self.wavelet,
                                               level=self.level,
                                               mode='periodization',
                                               axes=self.dirs),
                                 axes=(self.dirs))[0]
        return y.ravel() 

def __init__(self, dims, dir=0, wavelet='haar', level=1, dtype='float64'):
        if pywt is None:
            raise ModuleNotFoundError(pywt_message)
        _checkwavelet(wavelet)

        if isinstance(dims, int):
            dims = (dims, )

        # define padding for length to be power of 2
        ndimpow2 = max(2**ceil(log(dims[dir], 2)), 2 ** level)
        pad = [(0, 0)] * len(dims)
        pad[dir] = (0, ndimpow2 - dims[dir])
        self.pad = Pad(dims, pad)
        self.dims = dims
        self.dir = dir
        self.dimsd = list(dims)
        self.dimsd[self.dir] = ndimpow2

        # apply transform to find out slices
        _, self.sl = \
            pywt.coeffs_to_array(pywt.wavedecn(np.ones(self.dimsd),
                                               wavelet=wavelet,
                                               level=level,
                                               mode='periodization',
                                               axes=(self.dir,)),
                                 axes=(self.dir,))

        self.wavelet = wavelet
        self.waveletadj = _adjointwavelet(wavelet)
        self.level = level
        self.reshape = True if len(self.dims) > 1 else False
        self.shape = (int(np.prod(self.dimsd)), int(np.prod(self.dims)))
        self.dtype = np.dtype(dtype)
        self.explicit = False 

def _matvec(self, x):
        x = self.pad.matvec(x)
        if self.reshape:
            x = np.reshape(x, self.dimsd)
        y = pywt.coeffs_to_array(pywt.wavedecn(x, wavelet=self.wavelet,
                                               level=self.level,
                                               mode='periodization',
                                               axes=(self.dir,)),
                                 axes=(self.dir,))[0]
        return y.ravel() 

def wavelet_transform(x):
    w_coeffs_rgb = [] # np.zeros(x.shape[3], np.prod(x.shape))
    for i in range(x.shape[3]):
        w_coeffs_list = pywt.wavedec2(x[0,:,:,i], 'db4', level=None, mode='periodization')
        w_coeffs, coeff_slices = pywt.coeffs_to_array(w_coeffs_list)
        w_coeffs_rgb.append(w_coeffs)

    w_coeffs_rgb = np.array(w_coeffs_rgb)
    return w_coeffs_rgb, coeff_slices 

def test_wavedecn_coeff_reshape_axes_subset():
    # verify round trip is correct when only a subset of axes are transformed:
    #   wavedecn - >coeffs_to_array-> array_to_coeffs -> waverecn
    # This is done for wavedec{1, 2, n}
    rng = np.random.RandomState(1234)
    mode = 'symmetric'
    w = pywt.Wavelet('db2')
    N = 16
    ndim = 3
    for axes in [(-1, ), (0, ), (1, ), (0, 1), (1, 2), (0, 2), None]:
        x1 = rng.randn(*([N] * ndim))
        coeffs = pywt.wavedecn(x1, w, mode=mode, axes=axes)
        coeff_arr, coeff_slices = pywt.coeffs_to_array(coeffs, axes=axes)
        if axes is not None:
            # if axes is not None, it must be provided to coeffs_to_array
            assert_raises(ValueError, pywt.coeffs_to_array, coeffs)

        # mismatched axes size
        assert_raises(ValueError, pywt.coeffs_to_array, coeffs,
                      axes=(0, 1, 2, 3))
        assert_raises(ValueError, pywt.coeffs_to_array, coeffs,
                      axes=())

        coeffs2 = pywt.array_to_coeffs(coeff_arr, coeff_slices)
        x1r = pywt.waverecn(coeffs2, w, mode=mode, axes=axes)

        assert_allclose(x1, x1r, rtol=1e-4, atol=1e-4) 

def wavelet_transform(x):
    w_coeffs_rgb = [] # np.zeros(x.shape[3], np.prod(x.shape))
    for i in range(x.shape[3]):
        w_coeffs_list = pywt.wavedec2(x[0,:,:,i], 'db4', level=None, mode='periodization')
        w_coeffs, coeff_slices = pywt.coeffs_to_array(w_coeffs_list)
        w_coeffs_rgb.append(w_coeffs)

    w_coeffs_rgb = np.array(w_coeffs_rgb)
    return w_coeffs_rgb, coeff_slices 

def htrans(A):
    h0 = A
    res = []
    while h0.shape[0]>1 and h0.shape[1]>1:
        h0, (hx, hy, hc) = pywt.dwt2(h0, 'haar')
        res = [(h0, h0, h0)]+res
    out, _ = pywt.coeffs_to_array([h0]+res, padding=1)
    return out