NUFFT

Functional interface for the Non-Uniform Fast Fourier Transform.

torchlinops.functional.nufft

nufft(
    x: Tensor,
    locs: Float[Tensor, "... D"],
    oversamp: float = 1.25,
    width: float = 4.0,
)

Functional interface for the Non-Uniform Fast Fourier Transform.

Computes the forward NUFFT of input data at specified non-uniform locations. Internally applies apodization, zero-padding, FFT, and interpolation.

PARAMETER	DESCRIPTION
x	Input data on a regular grid. The last D dimensions are treated as spatial dimensions, where D = locs.shape[-1]. TYPE: Tensor
locs	Non-uniform sample locations. Each entry along the last dimension corresponds to a spatial axis and should lie in [-N//2, N//2] where N is the grid size along that axis. TYPE: Float[Tensor, '... D']
oversamp	Oversampling factor for the padded FFT grid. Default is 1.25. TYPE: float DEFAULT: 1.25
width	Interpolation kernel width. Default is 4.0. TYPE: float DEFAULT: 4.0

RETURNS	DESCRIPTION
Tensor	NUFFT values evaluated at the non-uniform locations.

Source code in src/torchlinops/functional/_nufft.py

def nufft(
    x: Tensor,
    locs: Float[Tensor, "... D"],
    oversamp: float = 1.25,
    width: float = 4.0,
):
    """Functional interface for the Non-Uniform Fast Fourier Transform.

    Computes the forward NUFFT of input data at specified non-uniform
    locations. Internally applies apodization, zero-padding, FFT, and
    interpolation.

    Parameters
    ----------
    x : Tensor
        Input data on a regular grid. The last ``D`` dimensions are
        treated as spatial dimensions, where ``D = locs.shape[-1]``.
    locs : Float[Tensor, "... D"]
        Non-uniform sample locations. Each entry along the last dimension
        corresponds to a spatial axis and should lie in
        ``[-N//2, N//2]`` where ``N`` is the grid size along that axis.
    oversamp : float, optional
        Oversampling factor for the padded FFT grid. Default is 1.25.
    width : float, optional
        Interpolation kernel width. Default is 4.0.

    Returns
    -------
    Tensor
        NUFFT values evaluated at the non-uniform locations.
    """

    grid_size = x.shape[-locs.shape[-1] :]
    params = init_nufft(grid_size, locs, oversamp, width, x.device)

    x = x * params.apodize
    x = Pad.fn(params.pad_ns, x)
    x = cfftn(x, dim=params.dim, norm="ortho")
    x = interpolate(
        x,
        params.locs,
        width,
        kernel="kaiser_bessel",
        kernel_params=dict(beta=params.beta),
    )
    x = x / params.scale_factor
    return x

torchlinops.functional.nufft_adjoint

nufft_adjoint(
    x: Tensor,
    locs: Float[Tensor, "... D"],
    grid_size: tuple[int, ...],
    oversamp: float = 1.25,
    width: float = 4.0,
)

Functional interface for the adjoint NUFFT.

Grids non-uniformly sampled data back onto a regular grid. Internally applies adjoint interpolation (gridding), inverse FFT, cropping, and apodization correction.

PARAMETER	DESCRIPTION
x	Non-uniformly sampled data to be gridded. TYPE: Tensor
locs	Non-uniform sample locations. Each entry along the last dimension corresponds to a spatial axis and should lie in [-N//2, N//2] where N is the grid size along that axis. TYPE: Float[Tensor, '... D']
grid_size	Desired output grid size for each spatial dimension. TYPE: tuple of int
oversamp	Oversampling factor for the padded FFT grid. Default is 1.25. TYPE: float DEFAULT: 1.25
width	Interpolation kernel width. Default is 4.0. TYPE: float DEFAULT: 4.0

RETURNS	DESCRIPTION
Tensor	Gridded data on a regular grid of shape (..., *grid_size).

Source code in src/torchlinops/functional/_nufft.py

def nufft_adjoint(
    x: Tensor,
    locs: Float[Tensor, "... D"],
    grid_size: tuple[int, ...],
    oversamp: float = 1.25,
    width: float = 4.0,
):
    """Functional interface for the adjoint NUFFT.

    Grids non-uniformly sampled data back onto a regular grid. Internally
    applies adjoint interpolation (gridding), inverse FFT, cropping, and
    apodization correction.

    Parameters
    ----------
    x : Tensor
        Non-uniformly sampled data to be gridded.
    locs : Float[Tensor, "... D"]
        Non-uniform sample locations. Each entry along the last dimension
        corresponds to a spatial axis and should lie in
        ``[-N//2, N//2]`` where ``N`` is the grid size along that axis.
    grid_size : tuple of int
        Desired output grid size for each spatial dimension.
    oversamp : float, optional
        Oversampling factor for the padded FFT grid. Default is 1.25.
    width : float, optional
        Interpolation kernel width. Default is 4.0.

    Returns
    -------
    Tensor
        Gridded data on a regular grid of shape ``(..., *grid_size)``.
    """
    params = init_nufft(grid_size, locs, oversamp, width, x.device)

    x = x / params.scale_factor
    x = interpolate_adjoint(
        x,
        params.locs,
        params.padded_size,
        width,
        kernel="kaiser_bessel",
        kernel_params=dict(beta=params.beta),
    )
    x = cifftn(x, dim=params.dim, norm="ortho")
    x = Pad.adj_fn(params.pad_ns, x)
    x = x * params.apodize
    return x