test_nfft.py

#!/usr/bin/env python3

from nfft import nfft
import matplotlib.pyplot as plt
import numpy as np
from numpy.fft import fft

def gen_values(times: np.ndarray) -> np.ndarray:
    "Compose three waves."
    freq = 1
    values = 3.0 * np.sin(2*np.pi*freq*times)
    freq = 4
    values += 1.0 * np.sin(2*np.pi*freq*times)
    freq = 7
    values += 0.5 * np.sin(2*np.pi*freq*times)
    return values


def plot_uniform_fft(count):
    times = np.arange(0,1,1.0/count)
    values = gen_values(times)

    spectrum = fft(values)
    freq = np.arange(count)

    n_oneside = count//2
    f_oneside = freq[:n_oneside]
    X_oneside = spectrum[:n_oneside]/n_oneside

    plt.subplot(231)
    plt.plot(times, values, 'r')
    plt.vlines(x=times, ymin=-2, ymax=2, colors='green', lw=.5)
    plt.ylabel('Uniform Samples')

    plt.subplot(232)
    plt.stem(freq, abs(spectrum), markerfmt=" ")
    plt.xlabel('Freq (Hz)')
    plt.ylabel('FFT')

    plt.subplot(233)
    plt.stem(f_oneside, abs(X_oneside), markerfmt=" ")
    plt.xlabel('Freq (Hz)')
    plt.ylabel('Normalized FFT')


def plot_nonuniform_fft(count):
    nutimes = np.linspace(0, 2*np.pi, count)
    nutimes = np.abs(np.sin(nutimes))
    nutimes = np.cumsum(nutimes)
    nutimes /= nutimes[-1]

    highres_domain = np.arange(0, 1, 1 / 256)
    highres_values = gen_values(highres_domain)
    nuvalues = np.interp(nutimes, highres_domain, highres_values)

    nuspectrum = nfft(nutimes, nuvalues)
    freq = np.arange(count)

    n_oneside = count//2
    f_oneside = freq[:n_oneside]
    X_oneside = nuspectrum[:n_oneside]/n_oneside

    plt.subplot(234)
    plt.plot(nutimes, nuvalues, "r")
    plt.vlines(x=nutimes, ymin=-2, ymax=2, colors="green", lw=.5)
    plt.ylabel("Non-Uniform Samples")

    plt.subplot(235)
    plt.stem(np.arange(count), abs(nuspectrum), markerfmt=" ")
    plt.xlabel('Freq (Hz)')
    plt.ylabel('NFFT')

    plt.subplot(236)
    plt.stem(f_oneside, abs(X_oneside), markerfmt=" ")
    plt.xlabel('Freq (Hz)')
    plt.ylabel('Normalized NFFT')


if __name__ == "__main__":
    plt.figure(figsize = (12, 6))
    plot_uniform_fft(64)
    plot_nonuniform_fft(64)
    plt.tight_layout()
    plt.show()