Skip to main content
GitHub Actions Documentation Status PyPI Versions PyPI - Project PyPI - License Black code style PyPi - downloads

Differentiable and GPU-enabled fast wavelet transforms in JAX.

Features

  • 1d analysis and synthesis transforms are implemented in src/jaxwt/conv_fwt.py.

  • 2d analysis and synthesis transform are part of the src/jaxwt/conv_fwt_2d.py module.

  • cwt-function supports 1d continuous wavelet transforms.

  • The WaveletPacket object supports 1d wavelet packet transforms.

  • WaveletPacket2d implements two-dimensional wavelet packet transforms.

Installation

To install Jax, head over to https://github.com/google/jax#installation and follow the procedure described there. Afterward, type pip install jaxwt to install the Jax-Wavelet-Toolbox. You can uninstall it later by typing pip uninstall jaxwt.

Documentation

The documentation is available at: https://jax-wavelet-toolbox.readthedocs.io .

Transform Examples:

One-dimensional fast wavelet transform:

import pywt
import numpy as np;
import jax.numpy as jnp
import jaxwt as jwt
# 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])
wavelet = pywt.Wavelet('haar')

# compare the forward fwt coefficients
print(pywt.wavedec(np.array(data), wavelet, mode='zero', level=2))
print(jwt.wavedec(data, wavelet, mode='zero', level=2))

# invert the fwt.
print(jwt.waverec(jwt.wavedec(data, wavelet, mode='zero', level=2),
                  wavelet))

Two-dimensional fast wavelet transform:

import pywt, scipy.misc
import jaxwt as jwt
import jax.numpy as jnp
face = jnp.transpose(
    scipy.misc.face(), [2, 0, 1]).astype(jnp.float64)
transformed = jwt.wavedec2(face, pywt.Wavelet("haar"),
                           level=2, mode="reflect")
reconstruction = jwt.waverec2(transformed, pywt.Wavelet("haar"))
jnp.max(jnp.abs(face - reconstruction))

Testing

Unit tests are handled by nox. Clone the repository and run it with the following:

$ pip install nox
$ git clone https://github.com/v0lta/Jax-Wavelet-Toolbox
$ cd Jax-Wavelet-Toolbox
$ nox -s test

Goals

  • In 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 numbers

If you need 64-bit floating point support, set the Jax config flag:

from jax.config import config
config.update("jax_enable_x64", True)

📖 Citation

If you use this work in a scientific context, please cite:

@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}
}

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

jaxwt-0.0.7.tar.gz (18.5 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

jaxwt-0.0.7-py3-none-any.whl (20.0 kB view details)

Uploaded Python 3

File details

Details for the file jaxwt-0.0.7.tar.gz.

File metadata

  • Download URL: jaxwt-0.0.7.tar.gz
  • Upload date:
  • Size: 18.5 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.1 CPython/3.10.4

File hashes

Hashes for jaxwt-0.0.7.tar.gz
Algorithm Hash digest
SHA256 d0d92f5ea6b93969dcb739378801549386b01b1ea7bd3e5e04d3a1e7305b3bc2
MD5 1f24e6eac5b410f7690310c68c164686
BLAKE2b-256 da7b0f7c28127f9c8f5d5ea8e2349ead4397505a1593b3bc6edd935b33493411

See more details on using hashes here.

File details

Details for the file jaxwt-0.0.7-py3-none-any.whl.

File metadata

  • Download URL: jaxwt-0.0.7-py3-none-any.whl
  • Upload date:
  • Size: 20.0 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.1 CPython/3.10.4

File hashes

Hashes for jaxwt-0.0.7-py3-none-any.whl
Algorithm Hash digest
SHA256 89372516979b915c084dad0c48037d7ee58e27eb690db8158e879fec3cbf06c3
MD5 bb62a955b8c92e4f2cb8e5b107ff553b
BLAKE2b-256 abec0ec9ebc5b738761a165a8dc86cc3891abe4cc550c34e89fd08b16fb2c409

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page