How to use the pylops.signalprocessing.FFT function in pylops

assert dottest(FFTop, nfft * nx * ny, nt * nx * ny,
                       complexflag=3, tol=10**(-decimal))

    D = FFTop * d.flatten()
    dadj = FFTop.H * D  # adjoint is same as inverse for fft
    dinv = lsqr(FFTop, D, damp=1e-10, iter_lim=10, show=0)[0]

    dadj = np.real(dadj.reshape(nt, nx, ny))
    dinv = np.real(dinv.reshape(nt, nx, ny))

    assert_array_almost_equal(d, dadj, decimal=decimal)
    assert_array_almost_equal(d, dinv, decimal=decimal)

    # 2nd dimension
    nfft = par['nx'] if par['nfft'] is None else par['nfft']
    FFTop = FFT(dims=(nt, nx, ny), dir=1, nfft=nfft, sampling=dt,
                real=par['real'], engine=par['engine'], dtype=par['dtype'])

    # FFT with real=True cannot pass dot-test neither be inverted correctly,
    # see FFT documentation for a detailed explanation. We thus test FFT.H*FFT
    if par['real']:
        FFTop = FFTop.H * FFTop
        assert dottest(FFTop, nt * nx * ny, nt * nx * ny,
                       complexflag=0, tol=10**(-decimal))
    else:
        assert dottest(FFTop, nt * nfft * ny, nt * nx * ny,
                       complexflag=2, tol=10**(-decimal))
        assert dottest(FFTop, nt * nfft * ny, nt * nx * ny,
                       complexflag=3, tol=10**(-decimal))

    D = FFTop * d.flatten()
    dadj = FFTop.H * D  # adjoint is inverse for fft

def test_FFT_1dsignal(par):
    """Dot-test and inversion for FFT operator for 1d signal
    """
    decimal = 4 if np.real(np.ones(1, par['dtype'])).dtype == np.float32 else 8

    dt = 0.005
    t = np.arange(par['nt']) * dt
    f0 = 10
    x = np.sin(2 * np.pi * f0 * t)
    x = x.astype(par['dtype'])
    nfft = par['nt'] if par['nfft'] is None else par['nfft']
    
    FFTop = FFT(dims=[par['nt']], nfft=nfft, sampling=dt,
                real=par['real'], engine=par['engine'], dtype=par['dtype'])

    # FFT with real=True cannot pass dot-test neither be inverted correctly,
    # see FFT documentation for a detailed explanation. We thus test FFT.H*FFT
    if par['real']:
        FFTop = FFTop.H * FFTop
        assert dottest(FFTop, par['nt'], par['nt'],
                       complexflag=0, tol=10**(-decimal), verb=True)
    else:
        assert dottest(FFTop, nfft, par['nt'],
                       complexflag=2, tol=10**(-decimal), verb=True)
        assert dottest(FFTop, nfft, par['nt'],
                       complexflag=3, tol=10**(-decimal), verb=True)

    y = FFTop * x
    xadj = FFTop.H * y  # adjoint is same as inverse for fft

def test_unknown_engine(par):
    """Check error is raised if unknown engine is passed
    """
    with pytest.raises(NotImplementedError):
        _ = FFT(dims=[par['nt']], nfft=par['nfft'], sampling=0.005,
                real=par['real'], engine='foo')

np.random.seed(10)

###############################################################################
# Let's first create the data in the frequency domain. The data is composed
# by the superposition of 3 sinusoids with different frequencies.

# Signal creation in frequency domain
ifreqs = [41, 25, 66]
amps = [1., 1., 1.]
N = 200
nfft = 2**11
dt = 0.004
t = np.arange(N)*dt
f = np.fft.rfftfreq(nfft, dt)

FFTop = 10*pylops.signalprocessing.FFT(N, nfft=nfft, real=True)

X = np.zeros(nfft//2+1, dtype='complex128')
X[ifreqs] = amps
x = FFTop.H*X

fig, axs = plt.subplots(2, 1, figsize=(12, 8))
axs[0].plot(f, np.abs(X), 'k', LineWidth=2)
axs[0].set_xlim(0, 30)
axs[0].set_title('Data(frequency domain)')
axs[1].plot(t, x, 'k', LineWidth=2)
axs[1].set_title('Data(time domain)')
axs[1].axis('tight')

###############################################################################
# We now define the locations at which the signal will be sampled.

phase
    """
    f = np.fft.rfftfreq(nfft, dt)
    df = f[1] - f[0]
    ifreqs = [int(np.random.normal(f, sigma)/df)
              for f, sigma in zip(f0, sigmaf)]
    amps = [np.random.normal(a, sigma) for a, sigma in zip(a0, sigmaa)]
    phis = [np.random.normal(phi, sigma) for phi, sigma in zip(phi0, sigmaphi)]

    # input signal in frequency domain
    X = np.zeros(nfft//2+1, dtype='complex128')
    X[ifreqs] = np.array(amps).squeeze() * \
                np.exp(1j * np.deg2rad(np.array(phis))).squeeze()

    # input signal in time domain
    FFTop = pylops.signalprocessing.FFT(nt, nfft=nfft, real=True)
    x = FFTop.H*X
    return x

axs[0].legend()
axs[0].set_title('Signal')
axs[1].plot(FFTop.f[:int(FFTop.nfft/2)],
            np.abs(D[:int(FFTop.nfft/2)]), 'k', lw=2)
axs[1].set_title('Fourier Transform')
axs[1].set_xlim([0, 3*f0])

###############################################################################
# In this example we used numpy as our engine for the ``fft`` and ``ifft``.
# PyLops implements a second engine (``engine='fftw'``) which uses the
# well-known `FFTW `_ via the python wrapper
# :py:class:`pyfftw.FFTW`. This optimized fft tends to outperform the one from
# numpy in many cases but it is not inserted in the mandatory requirements of
# PyLops. If interested to use ``FFTW`` backend, read the `fft routines`
# section at :ref:`performance`.
FFTop = pylops.signalprocessing.FFT(dims=nt, nfft=nfft,
                                    sampling=dt, engine='fftw')
D = FFTop * d

# Adjoint = inverse for FFT
dinv = FFTop.H * D
dinv = FFTop / D

fig, axs = plt.subplots(1, 2, figsize=(10, 4))
axs[0].plot(t, d, 'k', lw=2, label='True')
axs[0].plot(t, dinv, '--r', lw=2, label='Inverted')
axs[0].legend()
axs[0].set_title('Signal')
axs[1].plot(FFTop.f[:int(FFTop.nfft / 2)],
            np.abs(D[:int(FFTop.nfft / 2)]), 'k', lw=2)
axs[1].set_title('Fourier Transform with fftw')
axs[1].set_xlim([0, 3 * f0])

axs[1].plot(FFTop.f[:int(FFTop.nfft / 2)],
            np.abs(D[:int(FFTop.nfft / 2)]), 'k', lw=2)
axs[1].set_title('Fourier Transform with fftw')
axs[1].set_xlim([0, 3 * f0])

###############################################################################
# We can also apply the one dimensional FFT to to a two-dimensional
# signal (along one of the first axis)
dt = 0.005
nt, nx = 100, 20
t = np.arange(nt)*dt
f0 = 10
nfft = 2**10
d = np.outer(np.sin(2*np.pi*f0*t), np.arange(nx)+1)

FFTop = pylops.signalprocessing.FFT(dims=(nt, nx), dir=0,
                                    nfft=nfft, sampling=dt)
D = FFTop*d.flatten()

# Adjoint = inverse for FFT
dinv = FFTop.H*D
dinv = FFTop / D
dinv = np.real(dinv).reshape(nt, nx)

fig, axs = plt.subplots(2, 2, figsize=(10, 6))
axs[0][0].imshow(d, vmin=-20, vmax=20, cmap='seismic')
axs[0][0].set_title('Signal')
axs[0][0].axis('tight')
axs[0][1].imshow(np.abs(D.reshape(nfft, nx)[:200, :]), cmap='seismic')
axs[0][1].set_title('Fourier Transform')
axs[0][1].axis('tight')
axs[1][0].imshow(dinv, vmin=-20, vmax=20, cmap='seismic')

.. math::
        \mathbf{d} = \mathbf{F}^H_t \mathbf{F}^H_x  \mathbf{P}
            \mathbf{F}_x \mathbf{F}_t \mathbf{m}

    where :math:`\mathbf{P}` perfoms the phase-shift as discussed above.

    """
    dtypefft = (np.ones(1, dtype=dtype) + 1j * np.ones(1, dtype=dtype)).dtype
    if ky is None:
        dims = (nt, kx.size)
        dimsfft = (freq.size, kx.size)
    else:
        dims = (nt, kx.size, ky.size)
        dimsfft = (freq.size, kx.size, ky.size)
    Fop = FFT(dims, dir=0, nfft=nt, real=True, dtype=dtypefft)
    Kxop = FFT(dimsfft, dir=1, nfft=kx.size, real=False,
               fftshift=True, dtype=dtypefft)
    if ky is not None:
        Kyop = FFT(dimsfft, dir=2, nfft=ky.size, real=False,
                   fftshift=True, dtype=dtypefft)
    Pop = _PhaseShift(vel, dz, freq, kx, ky, dtypefft)
    if ky is None:
        Pop = Fop.H * Kxop * Pop * Kxop.H * Fop
    else:
        Pop = Fop.H * Kxop * Kyop * Pop * Kyop.H * Kxop.H * Fop
    return LinearOperator(Pop)

import pylops

plt.close('all')

###############################################################################
# Let's start by applying the one dimensional FFT to a one dimensional
# sinusoidal signal :math:`d(t)=sin(2 \pi f_0t)` using a time axis of
# lenght :math:`nt` and sampling :math:`dt`
dt = 0.005
nt = 100
t = np.arange(nt)*dt
f0 = 10
nfft = 2**10
d = np.sin(2*np.pi*f0*t)

FFTop = pylops.signalprocessing.FFT(dims=nt, nfft=nfft,
                                    sampling=dt, engine='numpy')
D = FFTop*d

# Adjoint = inverse for FFT
dinv = FFTop.H*D
dinv = FFTop / D

fig, axs = plt.subplots(1, 2, figsize=(10, 4))
axs[0].plot(t, d, 'k', lw=2, label='True')
axs[0].plot(t, dinv, '--r', lw=2, label='Inverted')
axs[0].legend()
axs[0].set_title('Signal')
axs[1].plot(FFTop.f[:int(FFTop.nfft/2)],
            np.abs(D[:int(FFTop.nfft/2)]), 'k', lw=2)
axs[1].set_title('Fourier Transform')
axs[1].set_xlim([0, 3*f0])

where :math:`\mathbf{P}` perfoms the phase-shift as discussed above.

    """
    dtypefft = (np.ones(1, dtype=dtype) + 1j * np.ones(1, dtype=dtype)).dtype
    if ky is None:
        dims = (nt, kx.size)
        dimsfft = (freq.size, kx.size)
    else:
        dims = (nt, kx.size, ky.size)
        dimsfft = (freq.size, kx.size, ky.size)
    Fop = FFT(dims, dir=0, nfft=nt, real=True, dtype=dtypefft)
    Kxop = FFT(dimsfft, dir=1, nfft=kx.size, real=False,
               fftshift=True, dtype=dtypefft)
    if ky is not None:
        Kyop = FFT(dimsfft, dir=2, nfft=ky.size, real=False,
                   fftshift=True, dtype=dtypefft)
    Pop = _PhaseShift(vel, dz, freq, kx, ky, dtypefft)
    if ky is None:
        Pop = Fop.H * Kxop * Pop * Kxop.H * Fop
    else:
        Pop = Fop.H * Kxop * Kyop * Pop * Kyop.H * Kxop.H * Fop
    return LinearOperator(Pop)