API Reference

flexft

CenteredDFT

Centered ordinary DFT operator.

Computes

\[ \operatorname{fftshift} \left( \operatorname{fft} \left( \operatorname{ifftshift}(\mathbf g) \right) \right). \]

Source code in src/flexft/core.py

class CenteredDFT:
    r"""Centered ordinary DFT operator.

    Computes

    $$
    \operatorname{fftshift}
    \left(
        \operatorname{fft}
        \left(
            \operatorname{ifftshift}(\mathbf g)
        \right)
    \right).
    $$
    """

    def __call__(self, g):
        g = jnp.asarray(g)
        if g.ndim != 1:
            raise ValueError(f"g must be one-dimensional, got shape {tuple(g.shape)}.")
        if g.shape[0] == 0:
            raise ValueError("g must contain at least one sample.")
        return fftshift(fft(ifftshift(g)))

CenteredFRDFT

Centered fractional discrete Fourier transform operator.

This wraps FRDFT with the phases needed to evaluate

\[ \sum_{n=0}^{N-1} \mathbf g[n] \exp\left[-i2\pi\alpha(m-c)(n-c)\right], \]

where \(c=\lfloor N/2\rfloor\).

Source code in src/flexft/core.py

class CenteredFRDFT:
    r"""Centered fractional discrete Fourier transform operator.

    This wraps [`FRDFT`][flexft.core.FRDFT] with the phases needed to evaluate

    $$
    \sum_{n=0}^{N-1}
    \mathbf g[n]
    \exp\left[-i2\pi\alpha(m-c)(n-c)\right],
    $$

    where $c=\lfloor N/2\rfloor$.
    """

    def __init__(self, N, alpha):
        self.N = _validate_positive_int(N)
        self.alpha = _validate_finite_scalar(alpha, name="alpha")

        n = np.arange(self.N, dtype=np.float64)
        c = self.N // 2
        self.pre = _unit_phase(self.alpha * c * n)
        self.post = _unit_phase(self.alpha * (c * n - c**2))
        self.frdft = FRDFT(self.N, self.alpha)

    def __call__(self, g):
        """Apply the centered fractional DFT to a vector of length ``N``."""
        g = _as_vector(g, length=self.N, name="g")
        return self.post * self.frdft(self.pre * g)

__call__

__call__(g)

Apply the centered fractional DFT to a vector of length N.

Source code in src/flexft/core.py

def __call__(self, g):
    """Apply the centered fractional DFT to a vector of length ``N``."""
    g = _as_vector(g, length=self.N, name="g")
    return self.post * self.frdft(self.pre * g)

FRDFT

Fractional discrete Fourier transform operator.

This class represents

\[ \operatorname{frdft}_{\alpha}(\mathbf g)[m] = \sum_{n=0}^{N-1} \mathbf g[n]\exp(-i2\pi\alpha mn). \]

The ordinary DFT is recovered when \(\alpha=1/N\). The transform is evaluated using the Bailey--Swarztrauber/Bluestein chirp-convolution algorithm. Constructing an FRDFT precomputes the chirps and convolution kernel so the operator can be reused with different vectors of the same length.

Parameters:

Name	Type	Description	Default
N		Length of the input and output vectors. Must be a positive integer.	required
alpha		Finite fractionality parameter \(\alpha\).	required

See Also

frdft One-shot convenience function.

Source code in src/flexft/core.py

class FRDFT:
    r"""Fractional discrete Fourier transform operator.

    This class represents

    $$
    \operatorname{frdft}_{\alpha}(\mathbf g)[m]
    =
    \sum_{n=0}^{N-1}
    \mathbf g[n]\exp(-i2\pi\alpha mn).
    $$

    The ordinary DFT is recovered when $\alpha=1/N$. The transform is evaluated
    using the Bailey--Swarztrauber/Bluestein chirp-convolution algorithm.
    Constructing an ``FRDFT`` precomputes the chirps and convolution kernel so the
    operator can be reused with different vectors of the same length.

    Parameters
    ----------
    N
        Length of the input and output vectors. Must be a positive integer.
    alpha
        Finite fractionality parameter $\alpha$.

    See Also
    --------
    frdft
        One-shot convenience function.
    """

    def __init__(self, N, alpha):
        self.N = _validate_positive_int(N)
        self.alpha = _validate_finite_scalar(alpha, name="alpha")
        # Float64 avoids int32 overflow in n**2 once n exceeds 46,340.
        n = np.arange(self.N, dtype=np.float64)
        # exp(i*pi*alpha*n**2) = exp(i*2*pi*(alpha/2)*n**2).
        theta = _unit_phase(0.5 * self.alpha * n**2)
        theta_negative = _unit_phase(0.5 * self.alpha * (n - self.N) ** 2)

        self.ThetaStar = jnp.conjugate(theta)
        self.Zfft = fft(jnp.concatenate((theta, theta_negative)))

    def __call__(self, g):
        """Apply the fractional DFT to a vector of length ``N``."""
        g = _as_vector(g, length=self.N, name="g")
        Y = jnp.pad(g * self.ThetaStar, (0, self.N))
        conv = ifft(fft(Y) * self.Zfft)[: self.N]
        return self.ThetaStar * conv

__call__

__call__(g)

Apply the fractional DFT to a vector of length N.

Source code in src/flexft/core.py

def __call__(self, g):
    """Apply the fractional DFT to a vector of length ``N``."""
    g = _as_vector(g, length=self.N, name="g")
    Y = jnp.pad(g * self.ThetaStar, (0, self.N))
    conv = ifft(fft(Y) * self.Zfft)[: self.N]
    return self.ThetaStar * conv

FlexFT

Approximate the continuous Fourier transform on uniform grids.

The transform convention is

\[ F(k) = \int f(x)\exp(-i2\pi kx)\,\mathrm dx. \]

The grids are

\[ \mathbf{x}[n] = x_0 + (n-c)\delta_x, \qquad \mathbf{k}[m] = k_0 + (m-c)\delta_k, \]

where \(c=\lfloor N/2\rfloor\). If dk is omitted, the FFT-compatible spacing \(\delta_k=1/(N\delta_x)\) is used.

Parameters:

Name	Type	Description	Default
N		Positive input and output length.	required
dx		Positive direct-space spacing \(\delta_x\).	required
dk		Positive reciprocal-space spacing \(\delta_k\). If omitted, use 1 / (N * dx).	None
x0		Direct-space grid centre. Defaults to zero.	0.0
k0		Reciprocal-space grid centre. Defaults to zero.	0.0

Source code in src/flexft/core.py

class FlexFT:
    r"""Approximate the continuous Fourier transform on uniform grids.

    The transform convention is

    $$
    F(k) = \int f(x)\exp(-i2\pi kx)\,\mathrm dx.
    $$

    The grids are

    $$
    \mathbf{x}[n] = x_0 + (n-c)\delta_x,
    \qquad
    \mathbf{k}[m] = k_0 + (m-c)\delta_k,
    $$

    where $c=\lfloor N/2\rfloor$. If ``dk`` is omitted, the FFT-compatible
    spacing $\delta_k=1/(N\delta_x)$ is used.

    Parameters
    ----------
    N
        Positive input and output length.
    dx
        Positive direct-space spacing $\delta_x$.
    dk
        Positive reciprocal-space spacing $\delta_k$. If omitted, use
        ``1 / (N * dx)``.
    x0
        Direct-space grid centre. Defaults to zero.
    k0
        Reciprocal-space grid centre. Defaults to zero.
    """

    def __init__(self, *, N, dx, dk=None, x0=0.0, k0=0.0):
        self.N = _validate_positive_int(N)
        self.dx = _validate_spacing(dx, name="dx")
        self.x0 = _validate_finite_scalar(x0, name="x0")
        self.k0 = _validate_finite_scalar(k0, name="k0")

        if dk is None:
            self.dk = 1.0 / (self.N * self.dx)
            self.core = CenteredDFT()
        else:
            self.dk = _validate_spacing(dk, name="dk")
            self.core = CenteredFRDFT(self.N, self.dx * self.dk)

        c = self.N // 2
        n = np.arange(self.N, dtype=np.float64)
        x = (n - c) * self.dx
        k = (n - c) * self.dk
        self.pre = 1.0 if self.k0 == 0 else _unit_phase(-self.k0 * x)
        self.post = 1.0 if self.x0 == 0 else _unit_phase(-self.x0 * (k + self.k0))

    def __call__(self, f):
        """Transform samples ``f`` with shape ``(N,)``."""
        f = _as_vector(f, length=self.N, name="f")
        return self.dx * self.post * self.core(self.pre * f)

__call__

__call__(f)

Transform samples f with shape (N,).

Source code in src/flexft/core.py

def __call__(self, f):
    """Transform samples ``f`` with shape ``(N,)``."""
    f = _as_vector(f, length=self.N, name="f")
    return self.dx * self.post * self.core(self.pre * f)

FlexFT2D

Approximate the 2D CFT on uniform tensor-product grids.

dx contains the direct-space spacings and is required. dk contains the optional reciprocal-space spacings. Each grid argument may be a scalar, which is applied to both axes, or an axis-specific pair. The transform is applied separably along axis 0 and then axis 1.

Source code in src/flexft/core.py

class FlexFT2D:
    """Approximate the 2D CFT on uniform tensor-product grids.

    ``dx`` contains the direct-space spacings and is required. ``dk`` contains
    the optional reciprocal-space spacings. Each grid argument may be a scalar,
    which is applied to both axes, or an axis-specific pair. The transform is
    applied separably along axis 0 and then axis 1.
    """

    def __init__(self, *, N, dx, dk=None, x0=0.0, k0=0.0):
        N1, N2 = _as_pair(N, name="N")
        dx1, dx2 = _as_pair(dx, name="dx")
        dk1, dk2 = _as_pair(dk, name="dk", allow_none=True)
        x01, x02 = _as_pair(x0, name="x0")
        k01, k02 = _as_pair(k0, name="k0")

        self.N = (
            _validate_positive_int(N1, name="N[0]"),
            _validate_positive_int(N2, name="N[1]"),
        )
        self.op1 = FlexFT(N=self.N[0], dx=dx1, dk=dk1, x0=x01, k0=k01)
        self.op2 = FlexFT(N=self.N[1], dx=dx2, dk=dk2, x0=x02, k0=k02)

        self.dx = (self.op1.dx, self.op2.dx)
        self.dk = (self.op1.dk, self.op2.dk)
        self.x0 = (self.op1.x0, self.op2.x0)
        self.k0 = (self.op1.k0, self.op2.k0)

        self._op1_vm = jax.vmap(self.op1, in_axes=1, out_axes=1)
        self._op2_vm = jax.vmap(self.op2, in_axes=0, out_axes=0)

    def __call__(self, f):
        f = _as_matrix(f, shape=self.N, name="f")
        return self._op2_vm(self._op1_vm(f))

IFlexFT

Approximate the inverse continuous Fourier transform on uniform grids.

This approximates

\[ f(x) = \int F(k)\exp(i2\pi kx)\,\mathrm dk. \]

dk always denotes reciprocal-space spacing and is required. dx always denotes direct-space spacing; when omitted, the FFT-compatible value 1 / (N * dk) is used.

Parameters:

Name	Type	Description	Default
N		Positive input and output length.	required
dk		Positive reciprocal-space spacing \(\delta_k\).	required
dx		Positive direct-space spacing \(\delta_x\). If omitted, use 1 / (N * dk).	None
x0		Direct-space grid centre. Defaults to zero.	0.0
k0		Reciprocal-space grid centre. Defaults to zero.	0.0

Source code in src/flexft/core.py

class IFlexFT:
    r"""Approximate the inverse continuous Fourier transform on uniform grids.

    This approximates

    $$
    f(x) = \int F(k)\exp(i2\pi kx)\,\mathrm dk.
    $$

    ``dk`` always denotes reciprocal-space spacing and is required. ``dx``
    always denotes direct-space spacing; when omitted, the FFT-compatible value
    ``1 / (N * dk)`` is used.

    Parameters
    ----------
    N
        Positive input and output length.
    dk
        Positive reciprocal-space spacing $\delta_k$.
    dx
        Positive direct-space spacing $\delta_x$. If omitted, use
        ``1 / (N * dk)``.
    x0
        Direct-space grid centre. Defaults to zero.
    k0
        Reciprocal-space grid centre. Defaults to zero.
    """

    def __init__(self, *, N, dk, dx=None, x0=0.0, k0=0.0):
        self.N = _validate_positive_int(N)
        self.dk = _validate_spacing(dk, name="dk")
        self.x0 = _validate_finite_scalar(x0, name="x0")
        self.k0 = _validate_finite_scalar(k0, name="k0")

        requested_dx = None if dx is None else _validate_spacing(dx, name="dx")
        self.forward_like = FlexFT(
            N=self.N,
            dx=self.dk,
            dk=requested_dx,
            x0=self.k0,
            k0=self.x0,
        )
        self.dx = self.forward_like.dk

    def __call__(self, F):
        """Inverse-transform samples ``F`` with shape ``(N,)``."""
        F = _as_vector(F, length=self.N, name="F")
        return jnp.conj(self.forward_like(jnp.conj(F)))

__call__

__call__(F)

Inverse-transform samples F with shape (N,).

Source code in src/flexft/core.py

def __call__(self, F):
    """Inverse-transform samples ``F`` with shape ``(N,)``."""
    F = _as_vector(F, length=self.N, name="F")
    return jnp.conj(self.forward_like(jnp.conj(F)))

IFlexFT2D

Approximate the inverse 2D CFT on uniform tensor-product grids.

dk contains the required reciprocal-space spacings. dx contains the optional direct-space spacings; if omitted, FFT-compatible values are used. Each grid argument may be a scalar, which is applied to both axes, or an axis-specific pair.

Source code in src/flexft/core.py

class IFlexFT2D:
    """Approximate the inverse 2D CFT on uniform tensor-product grids.

    ``dk`` contains the required reciprocal-space spacings. ``dx`` contains the
    optional direct-space spacings; if omitted, FFT-compatible values are used.
    Each grid argument may be a scalar, which is applied to both axes, or an
    axis-specific pair.
    """

    def __init__(self, *, N, dk, dx=None, x0=0.0, k0=0.0):
        N_pair = _as_pair(N, name="N")
        dk_pair = _as_pair(dk, name="dk")
        dx_pair = _as_pair(dx, name="dx", allow_none=True)
        x0_pair = _as_pair(x0, name="x0")
        k0_pair = _as_pair(k0, name="k0")

        self.forward_like = FlexFT2D(
            N=N_pair,
            dx=dk_pair,
            dk=dx_pair,
            x0=k0_pair,
            k0=x0_pair,
        )

        self.N = self.forward_like.N
        self.dk = self.forward_like.dx
        self.dx = self.forward_like.dk
        self.x0 = self.forward_like.k0
        self.k0 = self.forward_like.x0

    def __call__(self, F):
        F = _as_matrix(F, shape=self.N, name="F")
        return jnp.conj(self.forward_like(jnp.conj(F)))

flexft

flexft(f, *, dx, dk=None, x0=0.0, k0=0.0)

Apply a forward FlexFT without explicitly constructing a reusable plan.

Source code in src/flexft/core.py

def flexft(f, *, dx, dk=None, x0=0.0, k0=0.0):
    """Apply a forward FlexFT without explicitly constructing a reusable plan."""
    f = jnp.asarray(f)
    if f.ndim != 1:
        raise ValueError(f"f must be one-dimensional, got shape {tuple(f.shape)}.")
    return FlexFT(N=f.shape[0], dx=dx, dk=dk, x0=x0, k0=k0)(f)

flexft2d

flexft2d(f, *, dx, dk=None, x0=0.0, k0=0.0)

Apply a 2D forward FlexFT without constructing a reusable plan.

Source code in src/flexft/core.py

def flexft2d(f, *, dx, dk=None, x0=0.0, k0=0.0):
    """Apply a 2D forward FlexFT without constructing a reusable plan."""
    f = jnp.asarray(f)
    if f.ndim != 2:
        raise ValueError(f"f must be two-dimensional, got shape {tuple(f.shape)}.")
    return FlexFT2D(N=f.shape, dx=dx, dk=dk, x0=x0, k0=k0)(f)

frdft

frdft(g, alpha)

Apply a fractional DFT without explicitly constructing a reusable plan.

Source code in src/flexft/core.py

def frdft(g, alpha):
    """Apply a fractional DFT without explicitly constructing a reusable plan."""
    g = jnp.asarray(g)
    if g.ndim != 1:
        raise ValueError(f"g must be one-dimensional, got shape {tuple(g.shape)}.")
    return FRDFT(g.shape[0], alpha)(g)

iflexft

iflexft(F, *, dk, dx=None, x0=0.0, k0=0.0)

Apply an inverse FlexFT without constructing a reusable plan.

Source code in src/flexft/core.py

def iflexft(F, *, dk, dx=None, x0=0.0, k0=0.0):
    """Apply an inverse FlexFT without constructing a reusable plan."""
    F = jnp.asarray(F)
    if F.ndim != 1:
        raise ValueError(f"F must be one-dimensional, got shape {tuple(F.shape)}.")
    return IFlexFT(N=F.shape[0], dk=dk, dx=dx, x0=x0, k0=k0)(F)

iflexft2d

iflexft2d(F, *, dk, dx=None, x0=0.0, k0=0.0)

Apply a 2D inverse FlexFT without constructing a reusable plan.

Source code in src/flexft/core.py

def iflexft2d(F, *, dk, dx=None, x0=0.0, k0=0.0):
    """Apply a 2D inverse FlexFT without constructing a reusable plan."""
    F = jnp.asarray(F)
    if F.ndim != 2:
        raise ValueError(f"F must be two-dimensional, got shape {tuple(F.shape)}.")
    return IFlexFT2D(N=F.shape, dk=dk, dx=dx, x0=x0, k0=k0)(F)