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

This version

0.0.3

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

Uploaded Source

Built Distributions

jax_finufft-0.0.3-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.3-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.3-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.3-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.3-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.3-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.3-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.3-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.3-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.3-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.3.tar.gz.

File metadata

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

File hashes

Hashes for jax-finufft-0.0.3.tar.gz
Algorithm Hash digest
SHA256 9f401da936ef53c1549f8efe37cb623e81c8b93bc9fb5b2ef9c06b7f920cb0bc
MD5 d8580be04d860cc2010ce5eab8f35221
BLAKE2b-256 a57fb511040193939f6cf1ae426796ad4338728638ce86bf60743027264ab320

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for jax_finufft-0.0.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 8a0305217a557103087b42e7483ae925cd8bcebd5eb3219464a258696071327d
MD5 f4cb3e75131b78419a65f2469333c0f1
BLAKE2b-256 2d9c522c853f1e2a924d26c5d39123fe930b1f9412ce7a017290775ed41d2d0c

See more details on using hashes here.

File details

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

File metadata

  • Download URL: jax_finufft-0.0.3-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.7.1 importlib_metadata/4.8.2 pkginfo/1.8.2 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.9

File hashes

Hashes for jax_finufft-0.0.3-cp310-cp310-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 4140a0a3e1322afcb544892868466d4a9bce5d01c8b8d1de3c64bc2dfe7bba3e
MD5 60f93431cdb43c9655aa901f58e2d764
BLAKE2b-256 75d1651dcde7448483eaab8fdb053557d878993b90972f068110e4d1420aa204

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for jax_finufft-0.0.3-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 21430c030db57b3c4262ff3dbb8fdbd6d16e227a72baa169d28d8566d1489423
MD5 5696a08806214fb219bf6c578571af4d
BLAKE2b-256 eb8158dd603447f532b37da2d3d2b013f300b2cfb303ca1fcad35fe44a0036fc

See more details on using hashes here.

File details

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

File metadata

  • Download URL: jax_finufft-0.0.3-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.7.1 importlib_metadata/4.8.2 pkginfo/1.8.2 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.9

File hashes

Hashes for jax_finufft-0.0.3-cp39-cp39-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 6fdee8b31fb5a508f43ebb4f84d05d224e29f216e306cfc8140eaa6774be404d
MD5 15fc9a52267471ac38035904af61ca2e
BLAKE2b-256 783173fbd2dc2d92bd2142afa66fbf51e5b21346a69dbeaff6a26ddecd57e9aa

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for jax_finufft-0.0.3-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 c32a52deab24f898610214bff2b62135fd984d4da0c23a50c43e079eee919a6f
MD5 dfd47c90c385996facd133a14cddb11c
BLAKE2b-256 00770fe6db9f842fa63dd074d4d1d39d1172f4398dd34bb2f294dc7acf72698b

See more details on using hashes here.

File details

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

File metadata

  • Download URL: jax_finufft-0.0.3-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.7.1 importlib_metadata/4.8.2 pkginfo/1.8.2 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.9

File hashes

Hashes for jax_finufft-0.0.3-cp38-cp38-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 d32ab88f89a2ab361a5a515688b77edf4a4f14f736fcf0f4660a766136371074
MD5 99bb952204cd17f50b582feb2251b385
BLAKE2b-256 49df845d03860ed4c025e031de2fd9aeae7e1e48c39c5f7a856b022252bac412

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for jax_finufft-0.0.3-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 637e0bfc02ec1c4e8fb7ea7135387c6317f933b7fdf4957714e8cbbf005ad8eb
MD5 3ad04eea6e1bf3c1d3e232b992e20eba
BLAKE2b-256 272daaead94b01b56a0f8765c35b146e526e518a9d1a7bc037907bddd7f6e566

See more details on using hashes here.

File details

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

File metadata

  • Download URL: jax_finufft-0.0.3-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.7.1 importlib_metadata/4.8.2 pkginfo/1.8.2 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.9

File hashes

Hashes for jax_finufft-0.0.3-cp37-cp37m-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 4ad819f2ff5fef4f438edaa9fa93d4aeca919bb8188ee5ca18fe9c7a39a70fe0
MD5 cbed62db14d6f08c975a644f4337bf33
BLAKE2b-256 0e9739e52312617d9b8da9a0968832e9f32b496f8471fa658af9e26da4668203

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for jax_finufft-0.0.3-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 d6292294ac1955209a7709ba8e7b729017c33063af3abe7cef3ceeb2dfe9cf71
MD5 0c65423b1a55a739683e6fb0b9d28314
BLAKE2b-256 dc8013e4fc0749f3570285d2ed805f0ea7cb9362ac03851d0d5f2041319adc73

See more details on using hashes here.

File details

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

File metadata

  • Download URL: jax_finufft-0.0.3-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.7.1 importlib_metadata/4.8.2 pkginfo/1.8.2 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.9

File hashes

Hashes for jax_finufft-0.0.3-cp36-cp36m-macosx_10_9_x86_64.whl
Algorithm Hash digest
SHA256 4bc90b5d76734bb0e4be9f3df5694dafe82660699292df51fd5d5185cbca9c38
MD5 a129fbdb02f80b040adda69f6263ed1b
BLAKE2b-256 97296cc2b0f52e13e9e8cfda7be23421ff96d613fd9f018ef2158afcec06f4a6

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