Unofficial JAX bindings for finufft

Navigation

Verified details

These details have been verified by PyPI
Maintainers
Avatar for danfm from gravatar.com danfm Avatar for lehman.garrison from gravatar.com lehman.garrison

Unverified details

These details have not been verified by PyPI
Project links
Meta

Project description

JAX bindings to FINUFFT

This package provides a JAX interface to (a subset of) the Flatiron Institute Non-uniform Fast Fourier Transform (FINUFFT) library. Take a look at the FINUFFT docs for all the necessary definitions, conventions, and more information about the algorithms and their implementation. This package uses a low-level interface to directly expose the FINUFFT library to JAX's XLA backend, as well as implementing differentiation rules for the transforms.

Included features

This library is currently CPU-only, but GPU support is in the works using the cuFINUFFT library.

Type 1 and 2 transforms are supported in 1-, 2-, and 3-dimensions. All of these functions support forward, reverse, and higher-order differentiation, as well as batching using vmap.

Installation

For now, only a source build is supported.

For building, you should only need a recent version of Python (>3.6) and FFTW. At runtime, you'll need numpy, scipy, and jax. To set up such an environment, you can use conda (but you're welcome to use whatever workflow works for you!):

conda create -n jax-finufft -c conda-forge python=3.9 numpy scipy fftw
python -m pip install "jax[cpu]"

Then you can install from source using (don't forget the --recursive flag because FINUFFT is included as a submodule):

git clone --recursive https://github.com/dfm/jax-finufft
cd jax-finufft
python -m pip install .

Usage

This library provides two high-level functions (and these should be all that you generally need to interact with): nufft1 and nufft2 (for the two "types" of transforms). If you're already familiar with the Python interface to FINUFFT, please note that the function signatures here are different!

For example, here's how you can do a 1-dimensional type 1 transform:

import numpy as np
from jax_finufft import nufft1

M = 100000
N = 200000

x = 2 * np.pi * np.random.uniform(size=M)
c = np.random.standard_normal(size=M) + 1j * np.random.standard_normal(size=M)
f = nufft1(N, c, x, eps=1e-6, iflag=1)

Noting that the eps and iflag are optional, and that (for good reason, I promise!) the order of the positional arguments is reversed from the finufft Python package.

The syntax for a 2-, or 3-dimensional transform is:

f = nufft1((Nx, Ny), c, x, y)  # 2D
f = nufft1((Nx, Ny, Nz), c, x, y, z)  # 3D

The syntax for a type 2 transform is (also allowing optional iflag and eps parameters):

c = nufft2(f, x)  # 1D
c = nufft2(f, x, y)  # 2D
c = nufft2(f, x, y, z)  # 3D

Similar libraries

  • finufft: The "official" Python bindings to FINUFFT. A good choice if you're not already using JAX and if you don't need to differentiate through your transform.
  • mrphys/tensorflow-nufft: TensorFlow bindings for FINUFFT and cuFINUFFT.

License & attribution

This package, developed by Dan Foreman-Mackey is licensed under the Apache License, Version 2.0, with the following copyright:

Copyright 2021 The Simons Foundation, Inc.

If you use this software, please cite the primary references listed on the FINUFFT docs.

Project details

Verified details

These details have been verified by PyPI
Maintainers
Avatar for danfm from gravatar.com danfm Avatar for lehman.garrison from gravatar.com lehman.garrison

Unverified details

These details have not been verified by PyPI
Project links
Meta

Release history Release notifications | RSS feed

0.1.0

0.1.0rc2 pre-release

0.0.3

This version

0.0.2

0.0.1

0.0.1rc1 pre-release

Download files

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

Source Distribution

jax-finufft-0.0.2.tar.gz (4.7 MB view details)

Uploaded Source

Built Distributions

jax_finufft-0.0.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.6 MB view details)

Uploaded CPython 3.10 manylinux: glibc 2.17+ x86-64

jax_finufft-0.0.2-cp310-cp310-macosx_10_9_x86_64.whl (2.5 MB view details)

Uploaded CPython 3.10 macOS 10.9+ x86-64

jax_finufft-0.0.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.6 MB view details)

Uploaded CPython 3.9 manylinux: glibc 2.17+ x86-64

jax_finufft-0.0.2-cp39-cp39-macosx_10_9_x86_64.whl (2.5 MB view details)

Uploaded CPython 3.9 macOS 10.9+ x86-64

jax_finufft-0.0.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.6 MB view details)

Uploaded CPython 3.8 manylinux: glibc 2.17+ x86-64

jax_finufft-0.0.2-cp38-cp38-macosx_10_9_x86_64.whl (2.5 MB view details)

Uploaded CPython 3.8 macOS 10.9+ x86-64

jax_finufft-0.0.2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.6 MB view details)

Uploaded CPython 3.7m manylinux: glibc 2.17+ x86-64

jax_finufft-0.0.2-cp37-cp37m-macosx_10_9_x86_64.whl (2.5 MB view details)

Uploaded CPython 3.7m macOS 10.9+ x86-64

jax_finufft-0.0.2-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.6 MB view details)

Uploaded CPython 3.6m manylinux: glibc 2.17+ x86-64

jax_finufft-0.0.2-cp36-cp36m-macosx_10_9_x86_64.whl (2.5 MB view details)

Uploaded CPython 3.6m macOS 10.9+ x86-64

File details

Details for the file jax-finufft-0.0.2.tar.gz.

File metadata

  • Download URL: jax-finufft-0.0.2.tar.gz
  • Upload date:
  • Size: 4.7 MB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.6.0 importlib_metadata/4.8.2 pkginfo/1.7.1 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.8

File hashes

Hashes for jax-finufft-0.0.2.tar.gz
Algorithm Hash digest
SHA256 835827b13b54295585028570d30d67321863d1810dd37a735d38823fda020bf6
MD5 14cec19fb4b71e97363713321b0ff872
BLAKE2b-256 26dd72fd74fb62743b30a9bad44d2637aa3f1a6aad6e8ffa6a2ef6e53020f2b1

See more details on using hashes here.

File details

Details for the file jax_finufft-0.0.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for jax_finufft-0.0.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 f423a76c0179becc18c6cc7da9379bf1dae7f4e4004bd8b79eb78c66d854e9a9
MD5 a22bc428e94f9803d9f55214a2d9c15f
BLAKE2b-256 8bec8c26349b51f1e1e71edf8e04a9febb834dedddb2bc29adf007404866a36e

See more details on using hashes here.

File details

Details for the file jax_finufft-0.0.2-cp310-cp310-macosx_10_9_x86_64.whl.

File metadata

  • Download URL: jax_finufft-0.0.2-cp310-cp310-macosx_10_9_x86_64.whl
  • Upload date:
  • Size: 2.5 MB
  • Tags: CPython 3.10, macOS 10.9+ x86-64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.6.0 importlib_metadata/4.8.2 pkginfo/1.7.1 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.8

File hashes

Hashes for jax_finufft-0.0.2-cp310-cp310-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 89563e2fb53ea00013e5787991b36a7149ed6dd7a903878ec6fdf721d5c23c3c
MD5 c370575c7656a5c85d6b6162e150aaed
BLAKE2b-256 0e6bb932b76067f04353812111c66523fba46152d7a915a131a5e45e2b9c7b72

See more details on using hashes here.

File details

Details for the file jax_finufft-0.0.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for jax_finufft-0.0.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 87ba2092f2ee8ad7abaf4699860012054599a54b8aaa07469afbf45d32b88278
MD5 ee8e4f365343632f97a9f81a5a09eaf0
BLAKE2b-256 71c0071b861aaf6037e2c501ec9ee706c346b00e2d5ebe4cf13c9fd3bce6aeaa

See more details on using hashes here.

File details

Details for the file jax_finufft-0.0.2-cp39-cp39-macosx_10_9_x86_64.whl.

File metadata

  • Download URL: jax_finufft-0.0.2-cp39-cp39-macosx_10_9_x86_64.whl
  • Upload date:
  • Size: 2.5 MB
  • Tags: CPython 3.9, macOS 10.9+ x86-64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.6.0 importlib_metadata/4.8.2 pkginfo/1.7.1 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.8

File hashes

Hashes for jax_finufft-0.0.2-cp39-cp39-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 abae95cbfec95cdc9d18a3d433b3557387e7fc9621a8d9bfab7f381a184b392f
MD5 40da0a05a9d8e9867845630dcc1f2b57
BLAKE2b-256 a469a5a3d3e7d8e6e749e04594ecef102a00591e720e314df92c006e2caca078

See more details on using hashes here.

File details

Details for the file jax_finufft-0.0.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for jax_finufft-0.0.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 cd79c468ff337b71fc9b41b237a35e36457e58184ca1e11a537c756763d7f179
MD5 c1c529e3bf869bd0df6d7b25cf2e45b2
BLAKE2b-256 6b2d43b648577ec0b0676436e3fb6cd59e8c3d56b94471678ed4f4d8450fd05b

See more details on using hashes here.

File details

Details for the file jax_finufft-0.0.2-cp38-cp38-macosx_10_9_x86_64.whl.

File metadata

  • Download URL: jax_finufft-0.0.2-cp38-cp38-macosx_10_9_x86_64.whl
  • Upload date:
  • Size: 2.5 MB
  • Tags: CPython 3.8, macOS 10.9+ x86-64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.6.0 importlib_metadata/4.8.2 pkginfo/1.7.1 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.8

File hashes

Hashes for jax_finufft-0.0.2-cp38-cp38-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 9474e84f98314ebb80032255215d580dea9163450675ca762a102bc58769fe6f
MD5 405aba8eb8b8e9a423194bc98f392118
BLAKE2b-256 4178546834293fbe07beb5213c3396e4c9c50316ff0811a062a22434ecedd674

See more details on using hashes here.

File details

Details for the file jax_finufft-0.0.2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for jax_finufft-0.0.2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 0ed9ae958aabeda1898265a4d9f805c0afa0157018b5b8036049ddea0dc88de2
MD5 f983a9ff72b0a97c0a9c183ac0bcdbc9
BLAKE2b-256 128da8ccd4d2bc185a64639c92a12a78aad527838a13e33b634b0c76c2641e7e

See more details on using hashes here.

File details

Details for the file jax_finufft-0.0.2-cp37-cp37m-macosx_10_9_x86_64.whl.

File metadata

  • Download URL: jax_finufft-0.0.2-cp37-cp37m-macosx_10_9_x86_64.whl
  • Upload date:
  • Size: 2.5 MB
  • Tags: CPython 3.7m, macOS 10.9+ x86-64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.6.0 importlib_metadata/4.8.2 pkginfo/1.7.1 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.8

File hashes

Hashes for jax_finufft-0.0.2-cp37-cp37m-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 24c305fa56cde226740e7d21570f89d4654812c4574b6d65ebc286c14d780da1
MD5 c6b8c78b1960d1132d9b6660764c5d4c
BLAKE2b-256 aab63bdc99e00e2a3a115091c882b38882a638dfa71a296141b62b80e519be4d

See more details on using hashes here.

File details

Details for the file jax_finufft-0.0.2-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for jax_finufft-0.0.2-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 a52d83a63c43642a3ef4b2499422a4d925b5ea394fffd9445a7a5a1aa4cf00d2
MD5 1a5a05aa383d8879001f203d320f64a4
BLAKE2b-256 8e119661182195a7f3f62a6f3db54022cad76436885d52291e36f6935bdf42d3

See more details on using hashes here.

File details

Details for the file jax_finufft-0.0.2-cp36-cp36m-macosx_10_9_x86_64.whl.

File metadata

  • Download URL: jax_finufft-0.0.2-cp36-cp36m-macosx_10_9_x86_64.whl
  • Upload date:
  • Size: 2.5 MB
  • Tags: CPython 3.6m, macOS 10.9+ x86-64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.6.0 importlib_metadata/4.8.2 pkginfo/1.7.1 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.8

File hashes

Hashes for jax_finufft-0.0.2-cp36-cp36m-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 12127a6b84c347edb14f1ca9c5c2a8df73286de460a2b97bac08794c3c577fbc
MD5 807acf2351cede0306db78109a8e9e75
BLAKE2b-256 9046dca53412c0c17fa45bfd94770294c5cb106a136b02ce7fb9bc2c010d3846

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page