Package: jaxwt
Differentiable and GPU-enabled fast wavelet transforms in JAX.Featureswavedecandwaverecimplement 1d analysis and synthesis transforms.Similarly,wavedec2andwaverec2provide 2d transform support.Thecwt-function supports 1d continuous wavelet transforms.TheWaveletPacketobject supports 1d wavelet packet transforms.WaveletPacket2dimplements two-dimensional wavelet packet transforms.swtandiswtallow 1d-stationary transformations.This toolbox extendsPyWavelets.
We additionally provide GPU and gradient support via a Jax backend.InstallationTo install Jax, head over tohttps://github.com/google/jax#installationand follow the procedure described there.
Afterward, typepip install jaxwtto install the Jax-Wavelet-Toolbox. You can uninstall it later by typingpip uninstall jaxwt.DocumentationComplete documentation of all toolbox functions is available atreadthedocs.Transform Examples:To compute a one-dimensional fast wavelet transform, consider the code snippet below:importjax.numpyasjnpimportjaxwtasjwtimportpywtimportnumpyasnp;# generate an input of even length.data=jnp.array([0.,1,2,3,4,5,6,7,7,6,5,4,3,2,1,0])# compare the forward fwt coefficientsprint(pywt.wavedec(np.array(data),'haar',mode='zero',level=2))print(jwt.wavedec(data,'haar',mode='zero',level=2))# invert the fwt.print(jwt.waverec(jwt.wavedec(data,'haar',mode='zero',level=2),'haar'))The snipped also evaluates thepywtimplementation to demonstrate that the coefficients are the same.
Usejaxwtif you require gradient or GPU support.The process for two-dimensional fast wavelet transforms works similarly:importjaxwtasjwtimportjax.numpyasjnpfromscipy.datasetsimportfaceimage=jnp.transpose(face(),[2,0,1]).astype(jnp.float32)transformed=jwt.wavedec2(image,"haar",level=2,mode="reflect")reconstruction=jwt.waverec2(transformed,"haar")jnp.max(jnp.abs(image-reconstruction))jaxwtallows transforming batched data.
The example above moves the color channel to the front because wavedec2 transforms the last two axes by default.
We can avoid doing so by using theaxesargument. Consider the batched example below:importjaxwtasjwtimportjax.numpyasjnpfromscipy.datasetsimportfaceimage=jnp.stack([face(),face(),face()],axis=0).astype(jnp.float32)transformed=jwt.wavedec2(image,"haar",level=2,mode="reflect",axes=(1,2))reconstruction=jwt.waverec2(transformed,"haar",axes=(1,2))jnp.max(jnp.abs(image-reconstruction))For more code examples, follow the documentation link above or visit
theexamplesfolder.TestingUnit tests are handled bynox. Clone the repository and run it with the following:$pipinstallnox$gitclonehttps://github.com/v0lta/Jax-Wavelet-Toolbox$cdJax-Wavelet-Toolbox$nox-stestGoalsIn the spirit of Jax, the aim is to be 100% pywt compatible. Whenever possible, interfaces should be the same
results identical.64-Bit floating-point numbersIf you need 64-bit floating point support, set the Jax config flag:fromjax.configimportconfigconfig.update("jax_enable_x64",True)CitationIf you use this work in a scientific context, please cite the following:@phdthesis{handle:20.500.11811/9245,
  urn: https://nbn-resolving.org/urn:nbn:de:hbz:5-63361,
  author = {{Moritz Wolter}},
  title = {Frequency Domain Methods in Recurrent Neural Networks for Sequential Data Processing},
  school = {Rheinische Friedrich-Wilhelms-Universität Bonn},
  year = 2021,
  month = jul,
  url = {https://hdl.handle.net/20.500.11811/9245}
}