#!/usr/bin/env python
# coding: utf-8

# # Benchmarks
# 
# This notebook contains benchmarks against the [pynfft](https://pypi.python.org/pypi/pyNFFT) package, which is a Python wrapper of Fortran code

# In[1]:


import numpy as np


# In[2]:


# Import from main directory
import os, sys
sys.path.append(os.path.abspath('..'))

import nfft
nfft.__version__


# In[3]:


import pynfft
pynfft.__version__


# ## Benchmarking the Forward Transform

# Define some test data:

# In[4]:


def make_forward_data(M, N):
    x = -0.5 + np.random.rand(M)
    f_hat = np.random.randn(N) + 1j * np.random.randn(N)
    return x, f_hat


# Define a utility function around pynfft:

# In[5]:


def pynfft_forward(x, f_hat):
    M = len(x)
    N = len(f_hat)
    plan = pynfft.nfft.NFFT(N, M)
    plan.x = x
    plan.precompute()
    plan.f_hat = f_hat
    # Need copy because of bug in pynfft 1.x
    # See https://github.com/ghisvail/pyNFFT/issues/57
    return plan.trafo().copy()


# In[6]:


x, f_hat = make_forward_data(1000, 100000)


# In[7]:


out1 = nfft.nfft(x, f_hat)
out2 = pynfft_forward(x, f_hat)
np.allclose(out1, out2)


# In[8]:


get_ipython().run_line_magic('timeit', 'nfft.nfft(x, f_hat)')
get_ipython().run_line_magic('timeit', 'pynfft_forward(x, f_hat)')


# ## Benchmarking the Adjoint Transform

# Define some test data:

# In[9]:


def make_adjoint_data(M):
    x = -0.5 + np.random.rand(M)
    f = np.random.randn(M) + 1j * np.random.randn(M)
    return x, f


# Define a utility function around pynfft:

# In[10]:


def pynfft_adjoint(x, f, N):
    M = len(x)
    plan = pynfft.nfft.NFFT(N, M)
    plan.x = x
    plan.precompute()
    plan.f = f
    # Need copy because of bug in pynfft 1.x
    # See https://github.com/ghisvail/pyNFFT/issues/57
    return plan.adjoint().copy()


# In[11]:


x, f = make_adjoint_data(1000)
N = 100000


# In[12]:


out1 = nfft.nfft_adjoint(x, f, N)
out2 = pynfft_adjoint(x, f, N)

np.allclose(out1, out2)


# In[13]:


get_ipython().run_line_magic('timeit', 'nfft.nfft_adjoint(x, f, N, sigma=3)')
get_ipython().run_line_magic('timeit', 'pynfft_adjoint(x, f, N)')