Skip to main content
GitHub Actions PyPI Versions PyPI - Project PyPI - License Black code style

Differentiable and GPU enabled fast wavelet transforms in JAX.

Features

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

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

Installation

To install jax, head over to https://github.com/google/jax#installation and follow the procedure described there. Afterwards type pip install jaxwt to install the Jax-Wavelet-Toolbox.

Transform Example:

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))

Testing

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

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

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.3.tar.gz (14.2 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.3-py3-none-any.whl (23.4 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: jaxwt-0.0.3.tar.gz
  • Upload date:
  • Size: 14.2 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.7.0 importlib_metadata/4.8.2 pkginfo/1.8.2 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.8.5

File hashes

Hashes for jaxwt-0.0.3.tar.gz
Algorithm Hash digest
SHA256 ac9733a80e5f99220d4193a92d8d9b459e92094b99d2d900b42968517d678089
MD5 e619a1b9c64ab82d6f16ca0c69909b50
BLAKE2b-256 8318970de561b989f4187b1532c2c14ec626578f6d3efe7ad163fa4fc07f620a

See more details on using hashes here.

File details

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

File metadata

  • Download URL: jaxwt-0.0.3-py3-none-any.whl
  • Upload date:
  • Size: 23.4 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.7.0 importlib_metadata/4.8.2 pkginfo/1.8.2 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.8.5

File hashes

Hashes for jaxwt-0.0.3-py3-none-any.whl
Algorithm Hash digest
SHA256 122a6c01b9e726462624057dde7e80a4e8f438837a95b06a547e71cafd40fbb3
MD5 fc8fd7a7c371f4f3175968d5dab870e9
BLAKE2b-256 e212d434fd9f16af67a388e8ca6c8bd550275bd2f7cf89ce8f09e3e9924d2b69

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