# -*- coding: utf-8 -*-
"""Wavelet transform functions.
"""
import numpy as np
import pywt
from sigpy import backend, util
__all__ = ['fwt', 'iwt']
def get_wavelet_shape(shape, wave_name='db4', axes=None, level=None):
zshape = [((i + 1) // 2) * 2 for i in shape]
tmp = pywt.wavedecn(
np.zeros(zshape), wave_name, mode='zero', axes=axes, level=level)
tmp, coeff_slices = pywt.coeffs_to_array(tmp, axes=axes)
oshape = tmp.shape
return oshape, coeff_slices
[docs]def fwt(input, wave_name='db4', axes=None, level=None):
"""Forward wavelet transform.
Args:
input (array): Input array.
axes (None or tuple of int): Axes to perform wavelet transform.
wave_name (str): Wavelet name.
level (None or int): Number of wavelet levels.
"""
device = backend.get_device(input)
input = backend.to_device(input, backend.cpu_device)
zshape = [((i + 1) // 2) * 2 for i in input.shape]
zinput = util.resize(input, zshape)
coeffs = pywt.wavedecn(
zinput, wave_name, mode='zero', axes=axes, level=level)
output, _ = pywt.coeffs_to_array(coeffs, axes=axes)
output = backend.to_device(output, device)
return output
[docs]def iwt(input, oshape, coeff_slices, wave_name='db4', axes=None, level=None):
"""Inverse wavelet transform.
Args:
input (array): Input array.
oshape (tuple of ints): Output shape.
coeff_slices (list of slice): Slices to split coefficients.
axes (None or tuple of int): Axes to perform wavelet transform.
wave_name (str): Wavelet name.
level (None or int): Number of wavelet levels.
"""
device = backend.get_device(input)
input = backend.to_device(input, backend.cpu_device)
input = pywt.array_to_coeffs(input, coeff_slices, output_format='wavedecn')
output = pywt.waverecn(input, wave_name, mode='zero', axes=axes)
output = util.resize(output, oshape)
output = backend.to_device(output, device)
return output