tensor/signal/nufft.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package signal

import (
	"sourcedock.dev/petrbalvin/tensor/internal/base"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

import "math"

// The type-1 nonuniform fast Fourier transform. Interferometric
// imaging and irregularly sampled spectroscopy ask for the spectrum
// of data whose samples sit at arbitrary coordinates: f_k =
// Σ_j c_j·e^{2πi·k·x_j} over a uniform output grid, with the x_j
// spread anywhere in [−1/2, 1/2). The direct sum costs O(n·m); the
// gridding route costs O(n + m·log m) by scattering the samples onto
// an oversampled grid through a localised Gaussian kernel, running
// one FFT, and undoing the kernel's own Fourier footprint by pointwise
// division. The Gaussian is the kernel of choice because its Fourier
// image is Gaussian too, so the deconvolution is as local as the
// spreading.

// NUFFTType1 computes f_k = Σ_j c_j·e^{2πi·k·x_j} for k = 0 … n−1,
// where x holds the nonuniform sample coordinates in [−1/2, 1/2) and
// c the complex values. The answer carries the gridding error of the
// Gaussian kernel, a few digits short of the direct sum at the
// default kernel width, in exchange for the FFT's speed. Coordinates
// outside [−1/2, 1/2), mismatched lengths, or a non-positive output
// size is an error.
func NUFFTType1(x, c *core.Array, n int) (*core.Array, error) {
	const name = "NUFFTType1"
	if x.NDim() != 1 || c.NDim() != 1 {
		return nil, base.Errf("%s: coordinates and values must be vectors", name)
	}
	if x.Len() != c.Len() {
		return nil, base.Errf("%s: %d coordinates for %d values", name, x.Len(), c.Len())
	}
	if x.Dtype() == core.Complex {
		return nil, base.Errf("%s: the coordinates must be real, got %s", name, x.Dtype())
	}
	if n <= 0 {
		return nil, base.Errf("%s: the output size must be positive, got %d", name, n)
	}
	count := x.Len()
	// Oversampling factor and kernel geometry: the wider the kernel
	// relative to the grid, the smaller the aliasing and truncation
	// error, at linear cost in the spreading step.
	upsampled := 8
	for upsampled < 4*n {
		upsampled *= 2
	}
	const sigma = 2.0
	const radius = 10
	grid := make([]complex128, upsampled)
	coords := make([]float64, count)
	for j := range count {
		coords[j] = x.FloatAt(j)
		// NaN defeats both range comparisons below, so it is refused
		// by name: it would poison the whole grid and the answer
		// would come back all-NaN with no error.
		if math.IsNaN(coords[j]) {
			return nil, base.Errf("%s: coordinate %d is NaN", name, j)
		}
		if coords[j] < -0.5 || coords[j] >= 0.5 {
			return nil, base.Errf("%s: coordinate %d = %g lies outside [−1/2, 1/2)", name, j, coords[j])
		}
	}
	// Spread: each sample lands on the nearest grid point and bleeds
	// into its radius neighbours through the kernel. The sample's value is
	// read once, not once per tap.
	for j := range count {
		g := coords[j] * float64(upsampled)
		near := math.Floor(g + 0.5)
		frac := g - near
		origin := int(near)
		cj := c.ComplexAt(j)
		for o := -radius; o <= radius; o++ {
			idx := ((origin+o)%upsampled + upsampled) % upsampled
			t := float64(o) - frac
			grid[idx] += cj * complex(math.Exp(-t*t/(2*sigma*sigma)), 0)
		}
	}
	// One inverse FFT of the oversampled grid, scaled by 1/m exactly as
	// the IFFT entry point applies it: a division, not a multiplication
	// by the reciprocal, so the rounding matches.
	transform(grid, +1)
	m := complex(float64(upsampled), 0)
	for k := range upsampled {
		grid[k] /= m
	}
	// Undo the kernel: the discrete Fourier image of the sampled
	// Gaussian at frequency k, computed by the same short sum. The
	// Gaussian envelope does not depend on k, so it is built once.
	gaussian := make([]float64, 2*radius+1)
	for o := -radius; o <= radius; o++ {
		gaussian[o+radius] = math.Exp(-float64(o*o) / (2 * sigma * sigma))
	}
	out, oerr := core.Zeros(core.Complex, []int{n}...)
	if oerr != nil {
		return nil, oerr
	}
	spectrum := out.RawComplexes()
	for k := range n {
		ku := k
		if ku > upsampled/2 {
			ku = upsampled - ku
		}
		image := complex(0, 0)
		// The same taps in the same order; the row index is carried by
		// the range so the kernel weight needs no index add. The
		// rotation comes straight from math.Sincos, whose pair is the
		// one CmplxPolar multiplies through (verified bit-for-bit), so
		// each tap halves its trig work.
		for oi := range gaussian {
			o := oi - radius
			s, c := math.Sincos(-2 * math.Pi * float64(o) * float64(ku) / float64(upsampled))
			image += complex(gaussian[oi]*c, gaussian[oi]*s)
		}
		spectrum[k] = grid[k] / image * complex(float64(upsampled), 0)
	}
	return out, nil
}