Skip to content

Examples

All one-dimensional functions accept or return arrays of shape (N,). The grid parameters always have the same meaning:

  • dx: direct-space spacing
  • dk: reciprocal-space spacing
  • x0: direct-space grid centre
  • k0: reciprocal-space grid centre

One-shot forward transform

import jax.numpy as jnp
from flexft import flexft

N = 256
dx = 0.05
dk = 0.02
x = (jnp.arange(N) - N // 2) * dx
f = jnp.exp(-(x**2))

F = flexft(f, dx=dx, dk=dk)
k = (jnp.arange(N) - N // 2) * dk

Omit dk to use dk = 1 / (N * dx) and the centered ordinary DFT path.

Reusing a transform plan

Constructing FlexFT precomputes grid phases and, for a non-FFT-compatible spacing, the FFT of the convolution kernel. Reuse the object when transforming many arrays on the same grids.

from flexft import FlexFT

transform = FlexFT(N=N, dx=dx, dk=dk)
F1 = transform(f)
F2 = transform(2 * f)

Inverse transform

The inverse transform always requires dk. It optionally accepts dx; if it is omitted, dx = 1 / (N * dk) is used.

from flexft import iflexft

f_inverse_approximation = iflexft(F, dk=dk, dx=dx)

For independently chosen spacings, forward and inverse calls are quadrature approximations to their respective continuous transforms. They are not generally exact matrix inverses. The FFT-compatible spacing makes the paired discrete operations exact inverses up to floating-point roundoff.

Shifted grids

x0 = 1.5
k0 = -0.25
x = x0 + (jnp.arange(N) - N // 2) * dx
f = jnp.exp(-((x - x0) ** 2))

F = flexft(f, dx=dx, dk=dk, x0=x0, k0=k0)
k = k0 + (jnp.arange(N) - N // 2) * dk

Two-dimensional transforms

from flexft import flexft2d, iflexft2d

shape = (128, 96)
dx2 = (0.05, 0.08)
dk2 = (0.02, 0.03)

x1 = (jnp.arange(shape[0]) - shape[0] // 2) * dx2[0]
x2 = (jnp.arange(shape[1]) - shape[1] // 2) * dx2[1]
f2 = jnp.exp(-(x1[:, None] ** 2 + x2[None, :] ** 2))

F2 = flexft2d(f2, dx=dx2, dk=dk2)
f2_inverse_approximation = iflexft2d(F2, dk=dk2, dx=dx2)

Each 2D grid argument may be a pair ordered by array axis, as above, or a scalar that is applied equally to both axes. For example, a square isotropic transform can be written as:

F2 = flexft2d(f2_square, dx=0.05, dk=0.02)

Precision

FlexFT constructs its plan phases using reduced float64 arithmetic before converting them to the precision selected by JAX. To run transforms in 64-bit precision, enable it before constructing a plan:

import jax

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