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

0.0.2

This version

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.1.tar.gz (4.7 MB view details)

Uploaded Source

Built Distributions

jax_finufft-0.0.1-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.1-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.1-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.1-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.1-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.1-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.1-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.1-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.1-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.1-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.1.tar.gz.

File metadata

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

File hashes

Hashes for jax-finufft-0.0.1.tar.gz
Algorithm Hash digest
SHA256 bfddf2d3a3d63cd2c90c1f2ddd67f3402d2d9a92037df7bf4b1a1564dc87ffd9
MD5 52249285cb9292298fa47e5d08f50e26
BLAKE2b-256 5efe1fed1f5e566580a7d0f1d98e4c3155b9fae7b3e87f2d7c7de613c048a216

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for jax_finufft-0.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 16e67326ec9df831519e512a3be1cca72d400ade4493a4dfda8ff0e2d1c1a74c
MD5 c5ffb528b0388e11b1f5071888a8aee8
BLAKE2b-256 eefc8acc9261c28358a730488d07259e3fbc7d50ab4154ca4b6d1336b74f23b3

See more details on using hashes here.

File details

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

File metadata

  • Download URL: jax_finufft-0.0.1-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.5.0 importlib_metadata/4.8.1 pkginfo/1.7.1 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.7

File hashes

Hashes for jax_finufft-0.0.1-cp310-cp310-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 29771bdad2fd81dad90444163e270d8db5808973dd197e92d416045fea4f7ed0
MD5 e77429010e8c4b370d25d33b77d450eb
BLAKE2b-256 bfb1f7c737b7b233f2f75b833aebaa3714c49ecde8be7259a9342e506e7b7b51

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for jax_finufft-0.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 7c7f2ea7d64e4611c3c7732a6893f121ffc8fc9e26f2ef89d788a0ce5009edf2
MD5 db43dde9e3da3284fc5d385c5424548f
BLAKE2b-256 7fa466859f05418aff387c4b78a48878e3a5904c792e033f33476d9f7793d405

See more details on using hashes here.

File details

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

File metadata

  • Download URL: jax_finufft-0.0.1-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.5.0 importlib_metadata/4.8.1 pkginfo/1.7.1 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.7

File hashes

Hashes for jax_finufft-0.0.1-cp39-cp39-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 f0a5c39476b95d560a2ef9c20c4fefe7c7d3c18009b5dd2542f19e1b6d13bf2e
MD5 1863c496bf3fa07c8eac6f6191c9c372
BLAKE2b-256 4eb070f97bb4612b541da08b5a1d82061de76bb90554026cb37ff348996ef9fe

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for jax_finufft-0.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 fccbc4d6d6f3b509cb12e40d97e525b3784a59266d392642c7b1cf9e48ead57f
MD5 2253838261e8733ae78abee523279255
BLAKE2b-256 40a33b317b1b673b3f5c72dae5f572640155437fd881c942f71341f5c8df8592

See more details on using hashes here.

File details

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

File metadata

  • Download URL: jax_finufft-0.0.1-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.5.0 importlib_metadata/4.8.1 pkginfo/1.7.1 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.7

File hashes

Hashes for jax_finufft-0.0.1-cp38-cp38-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 cb35b51689401f3890cd5028d791b7d4ad247972a213067dc54feeecc7672f99
MD5 bd14862f9eb13740572943e95cb8fc81
BLAKE2b-256 81b8c71aa6f94a1578754dcb7bcb589c1b5383de2a034802f3a5f5fb37ab85c1

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for jax_finufft-0.0.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 e0a7c409e3c931eca43592ef52c2c3003ef8515bb5a8c9d64b5c6c2e98e5f9ce
MD5 65f1d4c8911cb5fbedec428fcf5ce9ca
BLAKE2b-256 3e4fb81db1c53d1a3c9fc72ac923be9146f3577945853437b1a5b4f2afec9efb

See more details on using hashes here.

File details

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

File metadata

  • Download URL: jax_finufft-0.0.1-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.5.0 importlib_metadata/4.8.1 pkginfo/1.7.1 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.7

File hashes

Hashes for jax_finufft-0.0.1-cp37-cp37m-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 d34aac381097be00b534547b89ed2fd5cd255068d901a041da5502ab20634eb4
MD5 fdba2d285de2a379af4c369473a53e84
BLAKE2b-256 ae1c7ea3f6519bdcd4c620e39bc763857e2c2f719757c33548cfd1360c7a82cf

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for jax_finufft-0.0.1-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 bbaad4639d97a77216067019f996653f57d2c58e9ec7149af325a742ca8fad73
MD5 bc321e30c3b2630458c72a861e85f733
BLAKE2b-256 57ec6cf00b6de06a43f5712731e7cae6f6f0ba7b91e613c8c839a84d16329427

See more details on using hashes here.

File details

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

File metadata

  • Download URL: jax_finufft-0.0.1-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.5.0 importlib_metadata/4.8.1 pkginfo/1.7.1 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.7

File hashes

Hashes for jax_finufft-0.0.1-cp36-cp36m-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 8150f93b4321a8e0e0029f18e3e5bddb12fb6a31cee7e82ea77a8bb7ad69397b
MD5 208f8ce7fa66c8332404667e020fb57c
BLAKE2b-256 92a291ead4de50edb3cc536cab0a0fcd6b615d1647cae0a806ec519daa551202

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