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 analysis transforms.

  • WaveletPacket2d implements two-dimensional wavelet packet forward 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 np
face = np.transpose(scipy.misc.face(), [2, 0, 1]).astype(np.float64)
transformed = jwt.wavedec2(face, pywt.Wavelet("haar"), level=2, mode="reflect")
reconstruction = jwt.waverec2(transformed, pywt.Wavelet("haar"))
np.max(np.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

To allow 64-bit precision numbers, a Jax config flag must be set as shown below:

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

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.6.tar.gz (17.3 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.6-py3-none-any.whl (19.1 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: jaxwt-0.0.6.tar.gz
  • Upload date:
  • Size: 17.3 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.6.tar.gz
Algorithm Hash digest
SHA256 2b8537d894b4c95e67a83f208f366026290d465e9760208ce30793a5356fe2da
MD5 38a0510d4d926d4b9e346c41cdf788ef
BLAKE2b-256 5a487280c7f858dd85946dad755367fe420ead86c42d882590782677bb59b303

See more details on using hashes here.

File details

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

File metadata

  • Download URL: jaxwt-0.0.6-py3-none-any.whl
  • Upload date:
  • Size: 19.1 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.6-py3-none-any.whl
Algorithm Hash digest
SHA256 8c53b12fc151c14ea8e09dc988554fb96083d0e94147f09ec253fd71db91f4a8
MD5 70bf3bcc014ebf6aadd1e4e670deed56
BLAKE2b-256 33cfb8cba2066cdce59cb709521924d540d6b653dad430cdeeae03c856078786

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